refactor: optimize remove and sorting methods in Heap and HeapAsync c…

…lasses for improved performance
ignlg · Dec 28, 2024 · f5550ae · f5550ae
1 parent 5c58a96
commit f5550ae
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 90 deletions.
diff --git a/src/Heap.ts b/src/Heap.ts
@@ -614,24 +614,28 @@ export class Heap<T> implements Iterable<T> {
    * @return {Boolean}      True if the heap was modified
    */
   remove(o?: T, callbackFn: IsEqual<T> = Heap.defaultIsEqual): boolean {
-    if (this.length > 0) {
-      if (o === undefined) {
-        this.pop();
-        return true;
-      } else {
-        const idx = this.indexOf(o, callbackFn);
-        if (idx >= 0) {
-          if (idx === 0) {
-            this.pop();
-          } else if (idx === this.length - 1) {
-            this.heapArray.pop();
-          } else {
-            this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
-            this._sortNodeUp(idx);
-            this._sortNodeDown(idx);
-          }
-          return true;
+    if (!this.heapArray.length) return false;
+    if (o === undefined) {
+      this.pop();
+      return true;
+    }
+    const queue = [0];
+    while (queue.length) {
+      const idx = queue.shift() as number;
+      if (callbackFn(this.heapArray[idx], o)) {
+        if (idx === 0) {
+          this.pop();
+        } else if (idx === this.heapArray.length - 1) {
+          this.heapArray.pop();
+        } else {
+          this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
+          this._sortNodeUp(idx);
+          this._sortNodeDown(idx);
         }
+        return true;
+      } else if (this.compare(this.heapArray[idx], o) <= 0) {
+        const children = Heap.getChildrenIndexOf(idx).filter((c) => c < this.heapArray.length);
+        queue.push(...children);
       }
     }
     return false;
@@ -808,26 +812,20 @@ export class Heap<T> implements Iterable<T> {
    * @param  {Number} i Index of the node
    */
   _sortNodeDown(i: number): void {
-    let moveIt = i < this.heapArray.length - 1;
-    const self = this.heapArray[i];
-
-    const getPotentialParent = (best: number, j: number) => {
-      if (this.heapArray.length > j && this.compare(this.heapArray[j], this.heapArray[best]) < 0) {
-        best = j;
+    const { length } = this.heapArray;
+    while (true) {
+      const left = 2 * i + 1;
+      const right = left + 1;
+      let best = i;
+      if (left < length && this.compare(this.heapArray[left], this.heapArray[best]) < 0) {
+        best = left;
       }
-      return best;
-    };
-
-    while (moveIt) {
-      const childrenIdx = Heap.getChildrenIndexOf(i);
-      const bestChildIndex = childrenIdx.reduce(getPotentialParent, childrenIdx[0]);
-      const bestChild = this.heapArray[bestChildIndex];
-      if (typeof bestChild !== 'undefined' && this.compare(self, bestChild) > 0) {
-        this._moveNode(i, bestChildIndex);
-        i = bestChildIndex;
-      } else {
-        moveIt = false;
+      if (right < length && this.compare(this.heapArray[right], this.heapArray[best]) < 0) {
+        best = right;
       }
+      if (best === i) break;
+      this._moveNode(i, best);
+      i = best;
     }
   }
 
@@ -836,15 +834,12 @@ export class Heap<T> implements Iterable<T> {
    * @param  {Number} i Index of the node
    */
   _sortNodeUp(i: number): void {
-    let moveIt = i > 0;
-    while (moveIt) {
+    while (i > 0) {
       const pi = Heap.getParentIndexOf(i);
-      if (pi >= 0 && this.compare(this.heapArray[pi], this.heapArray[i]) > 0) {
+      if (pi >= 0 && this.compare(this.heapArray[i], this.heapArray[pi]) < 0) {
         this._moveNode(i, pi);
         i = pi;
-      } else {
-        moveIt = false;
-      }
+      } else break;
     }
   }
 

diff --git a/src/HeapAsync.ts b/src/HeapAsync.ts
@@ -537,30 +537,28 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
    * @return {Boolean}      True if the heap was modified
    */
   async remove(o?: T, fn: AsyncIsEqual<T> = HeapAsync.defaultIsEqual): Promise<boolean> {
-    if (this.length > 0) {
-      if (o === undefined) {
-        await this.pop();
+    if (!this.heapArray.length) return false;
+    if (o === undefined) {
+      await this.pop();
+      return true;
+    }
+    const queue = [0];
+    while (queue.length) {
+      const idx = queue.shift() as number;
+      if (await fn(this.heapArray[idx], o)) {
+        if (idx === 0) {
+          await this.pop();
+        } else if (idx === this.heapArray.length - 1) {
+          this.heapArray.pop();
+        } else {
+          this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
+          await this._sortNodeUp(idx);
+          await this._sortNodeDown(idx);
+        }
         return true;
       } else {
-        let idx = -1;
-        for (let i = 0; i < this.heapArray.length; ++i) {
-          if (await fn(this.heapArray[i], o)) {
-            idx = i;
-            break;
-          }
-        }
-        if (idx >= 0) {
-          if (idx === 0) {
-            await this.pop();
-          } else if (idx === this.length - 1) {
-            this.heapArray.pop();
-          } else {
-            this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
-            await this._sortNodeUp(idx);
-            await this._sortNodeDown(idx);
-          }
-          return true;
-        }
+        const children = HeapAsync.getChildrenIndexOf(idx).filter((c) => c < this.heapArray.length);
+        queue.push(...children);
       }
     }
     return false;
@@ -737,29 +735,20 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
    * @param  {Number} i Index of the node
    */
   async _sortNodeDown(i: number): Promise<void> {
-    let moveIt = i < this.heapArray.length - 1;
-    const self = this.heapArray[i];
-
-    const getPotentialParent = async (best: number, j: number) => {
-      if (this.heapArray.length > j && (await this.compare(this.heapArray[j], this.heapArray[best])) < 0) {
-        best = j;
+    const length = this.heapArray.length;
+    while (true) {
+      const left = 2 * i + 1;
+      const right = left + 1;
+      let best = i;
+      if (left < length && (await this.compare(this.heapArray[left], this.heapArray[best])) < 0) {
+        best = left;
       }
-      return best;
-    };
-
-    while (moveIt) {
-      const childrenIdx = HeapAsync.getChildrenIndexOf(i);
-      let bestChildIndex = childrenIdx[0];
-      for (let j = 1; j < childrenIdx.length; ++j) {
-        bestChildIndex = await getPotentialParent(bestChildIndex, childrenIdx[j]);
-      }
-      const bestChild = this.heapArray[bestChildIndex];
-      if (typeof bestChild !== 'undefined' && (await this.compare(self, bestChild)) > 0) {
-        this._moveNode(i, bestChildIndex);
-        i = bestChildIndex;
-      } else {
-        moveIt = false;
+      if (right < length && (await this.compare(this.heapArray[right], this.heapArray[best])) < 0) {
+        best = right;
       }
+      if (best === i) break;
+      this._moveNode(i, best);
+      i = best;
     }
   }
 
@@ -768,15 +757,12 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
    * @param  {Number} i Index of the node
    */
   async _sortNodeUp(i: number): Promise<void> {
-    let moveIt = i > 0;
-    while (moveIt) {
+    while (i > 0) {
       const pi = HeapAsync.getParentIndexOf(i);
-      if (pi >= 0 && (await this.compare(this.heapArray[pi], this.heapArray[i])) > 0) {
+      if (pi >= 0 && (await this.compare(this.heapArray[i], this.heapArray[pi])) < 0) {
         this._moveNode(i, pi);
         i = pi;
-      } else {
-        moveIt = false;
-      }
+      } else break;
     }
   }