Index-tracked open-set heap: O(log n) decrease-key + sift-down fix

src/abstract/heap.ts

@@ -1,8 +1,18 @@
 import { PathData, PathNode } from './node'
 
+/**
+ * Min-heap (by node `f`) used as A*'s open set.
+ *
+ * The heap is 1-indexed (slot 0 is an unused sentinel) so a node at index `i`
+ * has children at `2i` and `2i + 1` and parent at `i >>> 1`.
+ *
+ * Every node remembers its own slot in `node.heapIdx`. That is the key detail:
+ * A* calls {@link update} a lot (decrease-key, whenever it finds a cheaper route
+ * to an already-open node), and tracking the index makes that O(log n) instead
+ * of an O(n) `indexOf` scan over the whole open set.
+ */
 export class BinaryHeapOpenSet<Data extends PathData, N extends PathNode<Data>> {
-  // Initialing the array heap and adding a dummy element at index 0
-  heap: N[] = [null] as any
+  private readonly heap: N[] = [null as unknown as N] // slot 0 is an unused sentinel
 
   size (): number {
     return this.heap.length - 1
@@ -13,67 +23,78 @@ export class BinaryHeapOpenSet<Data extends PathData, N extends PathNode<Data>>
   }
 
   push (val: N): void {
-    // Inserting the new node at the end of the heap array
     this.heap.push(val)
+    this.siftUp(this.heap.length - 1)
+  }
+
+  /**
+   * Restore the heap order after `val`'s `f` has DECREASED (A* found a cheaper
+   * path to it). `val.heapIdx` tells us exactly where it sits, so we only need
+   * to bubble it up — no search required.
+   */
+  update (val: N): void {
+    this.siftUp(val.heapIdx)
+  }
 
-    // Finding the correct position for the new node
-    let current = this.heap.length - 1
-    let parent = current >>> 1
+  pop (): N {
+    const heap = this.heap
+    const min = heap[1]
+    const last = heap.pop() as N // remove the final element
 
-    // Traversing up the parent node until the current node is greater than the parent
-    while (current > 1 && this.heap[parent].f > this.heap[current].f) {
-      [this.heap[parent], this.heap[current]] = [this.heap[current], this.heap[parent]]
-      current = parent
-      parent = current >>> 1
+    // If `min` was the only element, `last === min` and the heap is now empty.
+    if (heap.length > 1) {
+      heap[1] = last
+      last.heapIdx = 1
+      this.siftDown(1)
     }
+
+    min.heapIdx = -1 // no longer in the heap
+    return min
   }
 
-  update (val: N): void {
-    let current = this.heap.indexOf(val)
-    let parent = current >>> 1
-
-    // Traversing up the parent node until the current node is greater than the parent
-    while (current > 1 && this.heap[parent].f > this.heap[current].f) {
-      [this.heap[parent], this.heap[current]] = [this.heap[current], this.heap[parent]]
-      current = parent
-      parent = current >>> 1
+  /** Bubble the node at `i` toward the root until its parent is no larger. */
+  private siftUp (i: number): void {
+    const heap = this.heap
+    const node = heap[i]
+    const f = node.f
+
+    while (i > 1) {
+      const parentIdx = i >>> 1
+      const parent = heap[parentIdx]
+      if (parent.f <= f) break
+      heap[i] = parent
+      parent.heapIdx = i
+      i = parentIdx
     }
+
+    heap[i] = node
+    node.heapIdx = i
   }
 
-  pop (): N {
-    // Smallest element is at the index 1 in the heap array
-    const smallest = this.heap[1]
-
-    this.heap[1] = this.heap[this.heap.length - 1]
-    this.heap.splice(this.heap.length - 1)
-
-    const size = this.heap.length - 1
-
-    if (size < 2) return smallest
-
-    const val = this.heap[1]
-    let index = 1
-    let smallerChild = 2
-    const cost = val.f
-    do {
-      let smallerChildNode = this.heap[smallerChild]
-      if (smallerChild < size - 1) {
-        const rightChildNode = this.heap[smallerChild + 1]
-        if (smallerChildNode.f > rightChildNode.f) {
-          smallerChild++
-          smallerChildNode = rightChildNode
-        }
-      }
-      if (cost <= smallerChildNode.f) {
-        break
-      }
-      this.heap[index] = smallerChildNode
-      this.heap[smallerChild] = val
-      index = smallerChild
-
-      smallerChild *= 2
-    } while (smallerChild <= size)
-
-    return smallest
+  /** Push the node at `i` toward the leaves until both children are no smaller. */
+  private siftDown (i: number): void {
+    const heap = this.heap
+    const size = heap.length - 1
+    const node = heap[i]
+    const f = node.f
+
+    while (true) {
+      let child = i << 1
+      if (child > size) break
+
+      // Pick the smaller of the two children. `child < size` guarantees the
+      // right child (`child + 1`) is a real slot before we read it.
+      if (child < size && heap[child + 1].f < heap[child].f) child++
+
+      const childNode = heap[child]
+      if (f <= childNode.f) break
+
+      heap[i] = childNode
+      childNode.heapIdx = i
+      i = child
+    }
+
+    heap[i] = node
+    node.heapIdx = i
   }
 }

src/abstract/node.ts

@@ -9,6 +9,11 @@ export class PathNode<Data extends PathData> {
   g = 0
   h = 0
 
+  // Current position in the open-set heap. Maintained by BinaryHeapOpenSet so
+  // that decrease-key (update) is O(log n) instead of an O(n) indexOf scan.
+  // -1 means "not in the heap".
+  heapIdx = -1
+
   get f (): number {
     return this.g + this.h
   }
@@ -41,6 +46,7 @@ export class CPathNode<Data extends PathData> implements PathNode<Data> {
   g: number
   h: number
   f: number
+  heapIdx = -1
 
   update (g: number, h: number, data: Data | null, parent: PathNode<Data> | null): this {
     this.g = g

tests/heap.test.ts

@@ -0,0 +1,79 @@
+import assert from 'node:assert/strict'
+import test from 'node:test'
+
+import { BinaryHeapOpenSet } from '../src/abstract/heap'
+import { CPathNode } from '../src/abstract/node'
+
+function makeHeap () {
+  return new BinaryHeapOpenSet<any, CPathNode<any>>()
+}
+
+function drain (heap: ReturnType<typeof makeHeap>): number[] {
+  const out: number[] = []
+  while (!heap.isEmpty()) out.push(heap.pop().f)
+  return out
+}
+
+test('pops nodes in ascending f order', () => {
+  const heap = makeHeap()
+  const fs = [7, 1, 8, 2, 9, 5, 3, 4, 6, 0]
+  for (const f of fs) heap.push(new CPathNode(f, 0))
+  assert.deepEqual(drain(heap), [...fs].sort((a, b) => a - b))
+})
+
+test('pop keeps the minimum at the root even when the smaller child is the last slot', () => {
+  // Regression: the old sift-down used `smallerChild < size - 1`, which skipped
+  // the right child when it was the final slot, so this used to return [1, 8, ...].
+  const heap = makeHeap()
+  for (const f of [1, 8, 2, 9]) heap.push(new CPathNode(f, 0))
+  assert.deepEqual(drain(heap), [1, 2, 8, 9])
+})
+
+test('update() re-heaps after a decrease-key (index-based, no scan)', () => {
+  const heap = makeHeap()
+  const a = new CPathNode(5, 0)
+  const b = new CPathNode(4, 0)
+  const c = new CPathNode(3, 0)
+  heap.push(a)
+  heap.push(b)
+  heap.push(c)
+
+  a.update(0, 0, null, null) // a.f: 5 -> 0, now the smallest
+  heap.update(a)
+
+  assert.equal(heap.pop(), a) // a must surface first
+  assert.deepEqual(drain(heap), [3, 4])
+})
+
+test('size/isEmpty track the count and heapIdx clears on pop', () => {
+  const heap = makeHeap()
+  const n = new CPathNode(1, 0)
+  assert.equal(heap.isEmpty(), true)
+  heap.push(n)
+  assert.equal(heap.size(), 1)
+  assert.ok(n.heapIdx > 0)
+  assert.equal(heap.pop(), n)
+  assert.equal(heap.isEmpty(), true)
+  assert.equal(n.heapIdx, -1)
+})
+
+test('stress: random pushes and decrease-keys always pop in sorted order', () => {
+  const heap = makeHeap()
+  const nodes: Array<CPathNode<any>> = []
+  for (let i = 0; i < 200; i++) {
+    const node = new CPathNode(Math.floor(Math.random() * 1000), 0)
+    nodes.push(node)
+    heap.push(node)
+  }
+  // Decrease the key of a handful of still-open nodes.
+  for (let i = 0; i < 40; i++) {
+    const node = nodes[Math.floor(Math.random() * nodes.length)]
+    if (node.heapIdx === -1) continue
+    // Always a strict decrease-key (what A* does); update() only bubbles up.
+    node.update(node.g - 1 - Math.floor(Math.random() * 10), 0, null, null)
+    heap.update(node)
+  }
+  const out = drain(heap)
+  const sorted = [...out].sort((a, b) => a - b)
+  assert.deepEqual(out, sorted)
+})