Top-K Data Structure in Javascript

README.md

Below is a TypeScript implementation of a Top-K data structure that maintains the top K largest items inserted so far. Internally it uses a min-heap of size up to K. When a new item arrives:

• If the heap size is less than K, we push the new item.
• If the heap is at capacity (size == K) and the new item is greater than the minimum element in the heap, we pop the min and push the new item.
• Otherwise, we ignore the new item (it’s not among the top K).

The minimum of the heap is always at the root, which allows us to efficiently compare new items against the smallest item in the top K set.

min-heap.ts

/**
 * A generic min-heap implementation in TypeScript.
 */
export class MinHeap<T> {
  private data: T[];
  private comparator: (a: T, b: T) => number;

  /**
   * @param comparator A function that compares two items: 
   *   comparator(a, b) < 0 means a < b
   *   comparator(a, b) > 0 means a > b
   *   comparator(a, b) = 0 means a == b
   */
  constructor(comparator: (a: T, b: T) => number) {
    this.data = [];
    this.comparator = comparator;
  }

  /**
   * Returns the size of the heap.
   */
  public size(): number {
    return this.data.length;
  }

  /**
   * Returns the minimum element without removing it.
   */
  public peek(): T | undefined {
    return this.data[0];
  }

  /**
   * Inserts a new value into the heap.
   */
  public push(value: T): void {
    this.data.push(value);
    this.heapifyUp(this.data.length - 1);
  }

  /**
   * Removes and returns the minimum element (root).
   */
  public pop(): T | undefined {
    const size = this.data.length;
    if (size === 0) return undefined;

    // Swap root with the last element
    this.swap(0, size - 1);
    const minValue = this.data.pop();
    // Restore heap property
    if (this.data.length > 0) {
      this.heapifyDown(0);
    }
    return minValue;
  }

  /**
   * Moves an element at index `index` up to restore heap property.
   */
  private heapifyUp(index: number): void {
    let parentIndex = Math.floor((index - 1) / 2);
    while (
      index > 0 &&
      this.comparator(this.data[index], this.data[parentIndex]) < 0
    ) {
      this.swap(index, parentIndex);
      index = parentIndex;
      parentIndex = Math.floor((index - 1) / 2);
    }
  }

  /**
   * Moves an element at index `index` down to restore heap property.
   */
  private heapifyDown(index: number): void {
    const size = this.data.length;

    while (true) {
      let smallest = index; // assume current index is smallest
      const left = 2 * index + 1;
      const right = 2 * index + 2;

      // Check left child
      if (
        left < size &&
        this.comparator(this.data[left], this.data[smallest]) < 0
      ) {
        smallest = left;
      }

      // Check right child
      if (
        right < size &&
        this.comparator(this.data[right], this.data[smallest]) < 0
      ) {
        smallest = right;
      }

      // If no swaps needed, break
      if (smallest === index) break;

      // Otherwise swap and continue
      this.swap(index, smallest);
      index = smallest;
    }
  }

  /**
   * Swaps elements at two indices.
   */
  private swap(i: number, j: number): void {
    [this.data[i], this.data[j]] = [this.data[j], this.data[i]];
  }
}

MinHeap.md

1. Min-Heap Implementation in TypeScript

We first define a generic MinHeap class with the following operations:

• size(): Returns the number of elements in the heap.
• peek(): Returns the minimum (root) without removing it.
• push(value): Inserts a value, keeping the heap structure.
• pop(): Removes and returns the minimum (root).
• heapifyUp() and heapifyDown(): Internal methods to maintain heap order.

Top-K.md

2. TopK Data Structure

Now that we have a MinHeap, we can build a TopK structure that tracks the top ( K ) largest items inserted. When TopK is full (size == ( K )):

• If the new item is larger than the minimum in the heap, we pop the min and push the new item (since we found a better candidate for the top K).
• Otherwise, we do nothing.

top-k.ts

import { MinHeap } from './min-heap';

/**
 * Maintains the top K largest items.
 */
export class TopK<T> {
  private capacity: number;
  private heap: MinHeap<T>;
  
  /**
   * @param k The maximum number of items to store
   * @param comparator a function comparator(a, b) that 
   *        returns < 0 if a < b, 
   *                 = 0 if a == b,
   *                 > 0 if a > b.
   * 
   * By default, we want a min-heap, so if a < b => comparator(a, b) < 0
   * which means 'a' should be on top in a min-heap scenario.
   * But for top-K largest items, we store them in a min-heap,
   * so the smallest is always at the root => easy to pop out the smallest if needed.
   */
  constructor(k: number, comparator: (a: T, b: T) => number) {
    this.capacity = k;
    // MinHeap that orders items by 'comparator'
    this.heap = new MinHeap(comparator);
  }

  /**
   * Potentially add a new item to the top-K structure.
   */
  public add(item: T): void {
    // If we still have room, just push
    if (this.heap.size() < this.capacity) {
      this.heap.push(item);
      return;
    }

    // If the new item is bigger than the min of the heap => replace
    const min = this.heap.peek();
    if (min !== undefined && this.heapComparator(item, min) > 0) {
      // pop the smallest
      this.heap.pop();
      this.heap.push(item);
    }
  }

  /**
   * Return the current count of items in the top-K.
   */
  public size(): number {
    return this.heap.size();
  }

  /**
   * Returns a sorted array of the top-K items, largest to smallest.
   * 
   * NOTE: This is an O(K log K) operation, because we need to pop everything
   * from the heap if we want them strictly sorted. 
   * Alternatively, we could just return the internal data unsorted 
   * if an approximate top-K is enough for your usage.
   */
  public getTopK(): T[] {
    // Copy the heap data, then pop to retrieve items in ascending order
    // (since it's a min-heap).
    const items: T[] = [];
    const temp: T[] = [];

    // We'll clone the heap data (shallow) so we don't destroy the original
    const cloneData = [...(this as any).heap.data];
    const clone = new MinHeap<T>((a, b) => this.heapComparator(a, b));
    cloneData.forEach((el) => clone.push(el));

    // Now pop everything out in ascending order, 
    // and build 'items' array from largest to smallest
    while (clone.size() > 0) {
      temp.push(clone.pop() as T); 
    }
    // 'temp' is sorted ascending, so we reverse for largest -> smallest
    temp.reverse();
    items.push(...temp);
    return items;
  }

  /**
   * A private utility to re-use the comparator in a simpler way:
   * returns negative if a < b, zero if equal, positive if a > b.
   */
  private heapComparator(a: T, b: T) {
    return this.heap['comparator'](a, b);
  }
}