feat(Heap, HeapAsync): implement element eviction based on limit and add _worstIndex method

Enhance the add method to reject elements worse than the worst kept when the heap reaches its limit, and replace the worst element if a better one is added. New private method, _worstIndex, to find the index of the worst element for eviction. Update _applyLimit to remove the worst elements while maintaining heap properties. Add tests to validate the new functionality for both Heap and HeapAsync.

Should fix #656

Author: ignlg

src/Heap.ts

@@ -352,11 +352,21 @@ export class Heap<T> implements Iterable<T> {
    * Adds an element to the heap. Aliases: {@link offer}.
    * Same as: {@link push}(element).
    * @param {any} element Element to be added
-   * @return {Boolean} true
+   * @return {Boolean} true if added, false if limit exceeded and element not good enough
    */
   add(element: T): boolean {
+    if (this._limit > 0 && this.heapArray.length >= this._limit) {
+      const worstIdx = this._worstIndex();
+      if (this.compare(element, this.heapArray[worstIdx]) >= 0) {
+        return false; // New element is not better than worst keeper
+      }
+      // Replace worst with new element
+      this.heapArray[worstIdx] = element;
+      this._sortNodeUp(worstIdx);
+      this._sortNodeDown(worstIdx);
+      return true;
+    }
     this._sortNodeUp(this.heapArray.push(element) - 1);
-    this._applyLimit();
     return true;
   }
 
@@ -819,14 +829,21 @@ export class Heap<T> implements Iterable<T> {
   }
 
   /**
-   * Limit heap size if needed
+   * Limit heap size if needed, removing worst elements to keep best N
    */
   _applyLimit(): void {
     if (this._limit > 0 && this._limit < this.heapArray.length) {
       let rm = this.heapArray.length - this._limit;
-      // It's much faster than splice
-      while (rm) {
-        this.heapArray.pop();
+      while (rm > 0) {
+        const worstIdx = this._worstIndex();
+        // Swap with last and pop (standard heap removal for non-root)
+        if (worstIdx === this.heapArray.length - 1) {
+          this.heapArray.pop();
+        } else {
+          this.heapArray[worstIdx] = this.heapArray.pop() as T;
+          this._sortNodeUp(worstIdx);
+          this._sortNodeDown(worstIdx);
+        }
         --rm;
       }
     }
@@ -1032,6 +1049,23 @@ export class Heap<T> implements Iterable<T> {
     heap.init(list);
     return heap.peek();
   }
+
+  /**
+   * Find index of the worst element (for eviction when at limit).
+   * Worst is always among leaves (second half of array).
+   * @return {number} Index of worst element, -1 if empty
+   */
+  _worstIndex(): number {
+    if (this.heapArray.length === 0) return -1;
+    const start = this.heapArray.length >> 1; // First leaf
+    let worstIdx = start;
+    for (let i = start + 1; i < this.heapArray.length; i++) {
+      if (this.compare(this.heapArray[i], this.heapArray[worstIdx]) > 0) {
+        worstIdx = i;
+      }
+    }
+    return worstIdx;
+  }
 }
 
 export default Heap;

src/HeapAsync.ts

@@ -347,11 +347,21 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
    * Adds an element to the heap. Aliases: `offer`.
    * Same as: push(element)
    * @param {any} element Element to be added
-   * @return {Boolean} true
+   * @return {Boolean} true if added, false if limit exceeded and element not good enough
    */
   async add(element: T): Promise<boolean> {
+    if (this._limit > 0 && this.heapArray.length >= this._limit) {
+      const worstIdx = await this._worstIndex();
+      if ((await this.compare(element, this.heapArray[worstIdx])) >= 0) {
+        return false; // New element is not better than worst keeper
+      }
+      // Replace worst with new element
+      this.heapArray[worstIdx] = element;
+      await this._sortNodeUp(worstIdx);
+      await this._sortNodeDown(worstIdx);
+      return true;
+    }
     await this._sortNodeUp(this.heapArray.push(element) - 1);
-    this._applyLimit();
     return true;
   }
 
@@ -367,7 +377,7 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
     for (const l = this.length; i < l; ++i) {
       await this._sortNodeUp(i);
     }
-    this._applyLimit();
+    await this._applyLimit();
     return true;
   }
 
@@ -463,7 +473,7 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
     for (let i = HeapAsync.getParentIndexOf(this.length - 1); i >= 0; --i) {
       await this._sortNodeDown(i);
     }
-    this._applyLimit();
+    await this._applyLimit();
   }
 
   /**
@@ -502,12 +512,24 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
   }
 
   /**
-   * Set length limit of the heap.
-   * @return {Number}
+   * Set length limit of the heap without eviction.
+   * Note: Use setLimit() for async limit application with proper eviction.
+   * @param {Number} _l Limit value
    */
   set limit(_l: number) {
     this._limit = ~~_l;
-    this._applyLimit();
+    // Note: This setter cannot await _applyLimit(). Use setLimit() for async application.
+  }
+
+  /**
+   * Set length limit of the heap with async limit application.
+   * @param {Number} _l Limit value
+   * @return {Promise<number>} The limit value
+   */
+  async setLimit(_l: number): Promise<number> {
+    this._limit = ~~_l;
+    await this._applyLimit();
+    return this._limit;
   }
 
   /**
@@ -699,14 +721,21 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
   }
 
   /**
-   * Limit heap size if needed
+   * Limit heap size if needed, removing worst elements to keep best N
    */
-  _applyLimit(): void {
-    if (this._limit && this._limit < this.heapArray.length) {
+  async _applyLimit(): Promise<void> {
+    if (this._limit > 0 && this._limit < this.heapArray.length) {
       let rm = this.heapArray.length - this._limit;
-      // It's much faster than splice
-      while (rm) {
-        this.heapArray.pop();
+      while (rm > 0) {
+        const worstIdx = await this._worstIndex();
+        // Swap with last and pop (standard heap removal for non-root)
+        if (worstIdx === this.heapArray.length - 1) {
+          this.heapArray.pop();
+        } else {
+          this.heapArray[worstIdx] = this.heapArray.pop() as T;
+          await this._sortNodeUp(worstIdx);
+          await this._sortNodeDown(worstIdx);
+        }
         --rm;
       }
     }
@@ -914,6 +943,23 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
     await heap.init(list);
     return heap.peek();
   }
+
+  /**
+   * Find index of the worst element (for eviction when at limit).
+   * Worst is always among leaves (second half of array).
+   * @return {number} Index of worst element, -1 if empty
+   */
+  async _worstIndex(): Promise<number> {
+    if (this.heapArray.length === 0) return -1;
+    const start = this.heapArray.length >> 1; // First leaf
+    let worstIdx = start;
+    for (let i = start + 1; i < this.heapArray.length; i++) {
+      if ((await this.compare(this.heapArray[i], this.heapArray[worstIdx])) > 0) {
+        worstIdx = i;
+      }
+    }
+    return worstIdx;
+  }
 }
 
 export default HeapAsync;

tests/heap-async/heap-async-private-methods.test.ts

@@ -165,6 +165,11 @@ describe('HeapAsync private', function () {
           expect(await heap._topIdxOf(top)).toEqual(mid);
         });
       });
+      describe('#_worstIndex()', function () {
+        it('should return -1 for an empty heap', async function () {
+          expect(await heap._worstIndex()).toBe(-1);
+        });
+      });
     });
   });
 });

tests/heap-async/heap-async-public-methods.test.ts

@@ -218,10 +218,10 @@ describe('HeapAsync instances', function () {
       });
 
       describe('#limit', function () {
-        it('should limit the heap length', async function () {
+        it('should limit the heap length using setLimit', async function () {
           await heap.init(values);
           expect(heap.length).toEqual(values.length);
-          heap.limit = 5;
+          await heap.setLimit(5);
           expect(heap.limit).toEqual(5);
           expect(heap.length).toEqual(5);
           const otherValues = values.slice(0, Math.floor(values.length / 2));
@@ -231,6 +231,63 @@ describe('HeapAsync instances', function () {
         });
       });
 
+      describe('#limit keeps top N best values', function () {
+        it('should keep the N smallest values for min-heap', async function () {
+          const minHeap = new HeapAsync(HeapAsync.minComparatorNumber);
+          await minHeap.setLimit(3);
+          await minHeap.push(5, 1, 8, 2, 9, 3, 7);
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should keep the N largest values for max-heap', async function () {
+          const maxHeap = new HeapAsync(HeapAsync.maxComparatorNumber);
+          await maxHeap.setLimit(3);
+          await maxHeap.push(5, 1, 8, 2, 9, 3, 7);
+          expect(maxHeap.length).toEqual(3);
+          expect(maxHeap.toArray().sort((a, b) => b - a)).toEqual([9, 8, 7]);
+        });
+
+        it('should keep the N best values when limit is set after init', async function () {
+          const minHeap = new HeapAsync(HeapAsync.minComparatorNumber);
+          await minHeap.init([5, 1, 8, 2, 9, 3, 7]);
+          await minHeap.setLimit(3);
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should reject elements worse than the worst kept when at capacity', async function () {
+          const minHeap = new HeapAsync(HeapAsync.minComparatorNumber);
+          await minHeap.setLimit(3);
+          await minHeap.push(1, 2, 3);
+          expect(await minHeap.add(4)).toBe(false); // 4 is worse than 3
+          expect(await minHeap.add(5)).toBe(false); // 5 is worse than 3
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should accept elements better than the worst kept when at capacity', async function () {
+          const minHeap = new HeapAsync(HeapAsync.minComparatorNumber);
+          await minHeap.setLimit(3);
+          await minHeap.push(3, 4, 5);
+          expect(await minHeap.add(2)).toBe(true); // 2 is better than 5
+          expect(await minHeap.add(1)).toBe(true); // 1 is better than 4
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should maintain heap property after limit enforcement', async function () {
+          const minHeap = new HeapAsync(HeapAsync.minComparatorNumber);
+          await minHeap.setLimit(5);
+          for (let i = 100; i > 0; i--) {
+            await minHeap.push(i);
+          }
+          expect(minHeap.length).toEqual(5);
+          expect(await minHeap.check()).toBeUndefined();
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3, 4, 5]);
+        });
+      });
+
       describe('#peek()', function () {
         it('should return the top element of the heap', async function () {
           const minValue = top(...values);

tests/heap/heap-private-methods.test.ts

@@ -147,6 +147,11 @@ describe('Heap private', function () {
           expect(heap._topIdxOf(top)).toEqual(mid);
         });
       });
+      describe('#_worstIndex()', function () {
+        it('should return -1 for an empty heap', function () {
+          expect(heap._worstIndex()).toBe(-1);
+        });
+      });
     });
   });
 });

tests/heap/heap-public-methods.test.ts

@@ -421,6 +421,63 @@ describe('Heap instances', function () {
         });
       });
 
+      describe('#limit keeps top N best values', function () {
+        it('should keep the N smallest values for min-heap', function () {
+          const minHeap = new Heap(Heap.minComparatorNumber);
+          minHeap.limit = 3;
+          minHeap.push(5, 1, 8, 2, 9, 3, 7);
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should keep the N largest values for max-heap', function () {
+          const maxHeap = new Heap(Heap.maxComparatorNumber);
+          maxHeap.limit = 3;
+          maxHeap.push(5, 1, 8, 2, 9, 3, 7);
+          expect(maxHeap.length).toEqual(3);
+          expect(maxHeap.toArray().sort((a, b) => b - a)).toEqual([9, 8, 7]);
+        });
+
+        it('should keep the N best values when limit is set after init', function () {
+          const minHeap = new Heap(Heap.minComparatorNumber);
+          minHeap.init([5, 1, 8, 2, 9, 3, 7]);
+          minHeap.limit = 3;
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should reject elements worse than the worst kept when at capacity', function () {
+          const minHeap = new Heap(Heap.minComparatorNumber);
+          minHeap.limit = 3;
+          minHeap.push(1, 2, 3);
+          expect(minHeap.add(4)).toBe(false); // 4 is worse than 3 (the worst kept)
+          expect(minHeap.add(5)).toBe(false); // 5 is worse than 3
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should accept elements better than the worst kept when at capacity', function () {
+          const minHeap = new Heap(Heap.minComparatorNumber);
+          minHeap.limit = 3;
+          minHeap.push(3, 4, 5);
+          expect(minHeap.add(2)).toBe(true); // 2 is better than 5 (the worst kept)
+          expect(minHeap.add(1)).toBe(true); // 1 is better than 4
+          expect(minHeap.length).toEqual(3);
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3]);
+        });
+
+        it('should maintain heap property after limit enforcement', function () {
+          const minHeap = new Heap(Heap.minComparatorNumber);
+          minHeap.limit = 5;
+          for (let i = 100; i > 0; i--) {
+            minHeap.push(i);
+          }
+          expect(minHeap.length).toEqual(5);
+          expect(minHeap.check()).toBeUndefined();
+          expect(minHeap.toArray().sort((a, b) => a - b)).toEqual([1, 2, 3, 4, 5]);
+        });
+      });
+
       describe('#peek()', function () {
         it('should return the top element of the heap', function () {
           const minValue = top(...values);