TypeScript Data Structures

heap.ts

type HeapNodeCompareFn<T> = (a: T, b: T) => number;

class Heap<T> {
  private nodes: T[];
  private compare: HeapNodeCompareFn<T>;

  constructor(compareFn: HeapNodeCompareFn<T>, nodes: T[] = []) {
    this.compare = compareFn;
    this.nodes = nodes;
    if (this.nodes.length > 0) {
      const lastParentIndex = Math.trunc(this.nodes.length / 2 - 1);
      for (let i = 0; i <= lastParentIndex; i++) {
        this.heapifyDown(i);
      }
    }
  }

  get size(): number {
    return this.nodes.length;
  }

  public insert(node: T) {
    this.nodes.push(node);
    this.heapifyUp(this.lastNodeIndex());
  }

  public extract(): T | undefined {
    if (this.size === 0) return undefined;
    this.swapByIndex(0, this.lastNodeIndex());
    const extractedNode = this.nodes.pop();
    this.heapifyDown(0);
    return extractedNode;
  }

  public peek(): T | undefined {
    return this.nodes[0];
  }

  public toArray(): T[] {
    return this.nodes;
  }

  private heapifyDown(startIndex: number) {
    let currentIndex = startIndex;
    while (currentIndex * 2 + 1 < this.nodes.length) {
      let childIndexToSwap = this.getChildIndexToSwap(currentIndex);
      if (childIndexToSwap !== undefined && this.compareByIndex(childIndexToSwap, currentIndex) < 0) {
        this.swapByIndex(currentIndex, childIndexToSwap);
        currentIndex = childIndexToSwap;
      } else {
        break;
      }
    }
  }

  private getChildIndexToSwap(parentIndex: number): number | undefined {
    const leftIndex = (parentIndex * 2) + 1;
    const rightIndex = (parentIndex * 2) + 2;

    if (rightIndex < this.nodes.length) {
      return this.compareByIndex(leftIndex, rightIndex) < 0 ? leftIndex : rightIndex;
    } else if (leftIndex < this.nodes.length) {
      return leftIndex;
    }
    return undefined;
  }

  private heapifyUp(startIndex: number) {
    let currentIndex = startIndex;
    while (currentIndex > 0) {
      let parentIndex = this.getParentIndex(currentIndex);
      if (this.compareByIndex(currentIndex, parentIndex) < 0) {
        this.swapByIndex(currentIndex, parentIndex);
        currentIndex = parentIndex;
      } else {
        break;
      }
    }
  }

  private lastNodeIndex(): number {
    return this.nodes.length - 1;
  }

  private getParentIndex(childIndex: number): number {
    return Math.trunc((childIndex - 1) / 2);
  }

  private compareByIndex(aIndex: number, bIndex: number): number {
    return this.compare(this.nodes[aIndex], this.nodes[bIndex]);
  }

  private swapByIndex(aIndex: number, bIndex: number) {
    [this.nodes[aIndex], this.nodes[bIndex]] = [this.nodes[bIndex], this.nodes[aIndex]];
  }
}

function createNumericMaxHeap(values: number[] = []): Heap<number> {
  return new Heap<number>((a, b) => b - a, values);
}

function createNumericMinHeap(values: number[] = []): Heap<number> {
  return new Heap<number>((a, b) => a - b, values);
}


export { Heap, createNumericMaxHeap, createNumericMinHeap };