Heap | TypeScript/JavaScript Implementation for Technical Coding Interviews

00-heap.ts

/**
 * Heap implementation.
 *
 * Adapted from a Python implementation in the "Beyond Cracking the Coding Interview" book.
 */

type PriorityCompFunc <T> = (a: T, b: T) => boolean;

class Heap <T> {
  readonly heap: T[];
  readonly priorityComp: PriorityCompFunc <T>;

  constructor({
    priorityComp,
    heap,
  }: {
    priorityComp ? : PriorityCompFunc <T>,
    heap ? : T[]
  } = {}) {
    this.heap = [];
    if (heap !== undefined) this.heap = heap;

    this.priorityComp = (a: T, b: T) => a < b;
    if (priorityComp !== undefined) this.priorityComp = priorityComp;

    if (this.heap.length > 0) this.heapify();
  }

  /**
   * PUBLIC INTERFACE
   */
  push(elem: T): void {
    this.heap.push(elem);
    this.bubbleUp(this.heap.length - 1);
  }

  pop(): T | null {
    if (this.heap.length === 0) return null;

    const top = this.heap[0];
    if (this.heap.length === 1) {
      this.heap.pop();
      return top;
    }

    this.heap[0] = this.heap[this.heap.length - 1];
    this.heap.pop();
    this.bubbleDown(0);

    return top;
  }

  top(): T | null {
    if (this.heap.length === 0) return null;
    return this.heap[0];
  }

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

  /**
   * PRIVATE METHODS
   */
  /**
   * Bubble up all internal nodes. In a complete tree, at least half
   * the nodes are leaves, so we can start bubbling down from the middle
   * of the array.
   */
  private heapify(): void {
    for (let idx = Math.floor(this.heap.length / 2); idx >= 0; idx--) {
      this.bubbleDown(idx);
    }
  }

  private bubbleUp(idx: number): void {
    const parentIdx = this.parent(idx);
    // the root cannot be bubbled up.
    if (idx === 0 || parentIdx === null) return;

    if (this.priorityComp(this.heap[idx], this.heap[parentIdx])) {
      [this.heap[idx], this.heap[parentIdx]] = [
        this.heap[parentIdx],
        this.heap[idx],
      ];
      this.bubbleUp(parentIdx);
    }
  }

  private bubbleDown(idx: number): void {
    const [lIdx, rIdx] = [this.leftChild(idx), this.rightChild(idx)];
    const isLeaf = lIdx >= this.heap.length;

    // Leaves cannot be bubbled down.
    if (isLeaf) return;

    // Index for the highest priority child.
    let childIdx = lIdx;

    if (
      rIdx < this.heap.length &&
      this.priorityComp(this.heap[rIdx], this.heap[lIdx])
    ) {
      childIdx = rIdx;
    }

    if (this.priorityComp(this.heap[childIdx], this.heap[idx])) {
      [this.heap[idx], this.heap[childIdx]] = [
        this.heap[childIdx],
        this.heap[idx],
      ];
      this.bubbleDown(childIdx);
    }
  }

  private parent(idx: number) {
    // root has no parent.
    if (idx === 0) return null;

    return Math.floor((idx - 1) / 2);
  }

  private leftChild(idx: number) {
    return 2 * idx + 1;
  }

  private rightChild(idx: number) {
    return 2 * idx + 2;
  }
}

01-tests.ts

function runTests() {
  // Test min heap
  const minHeap = new Heap < number > ();
  const values = [4, 8, 2, 6, 1, 7, 3, 5];
  for (const val of values) {
    minHeap.push(val);
  }

  // Should pop in ascending order
  const sortedValues1 = [];
  while (minHeap.size() > 0) {
    sortedValues1.push(minHeap.pop());
  }
  if (
    JSON.stringify(sortedValues1) !== JSON.stringify([1, 2, 3, 4, 5, 6, 7, 8])
  ) {
    throw new Error(
      `\nmin_heap popped values: got: ${sortedValues1}, want: [1, 2, 3, 4, 5, 6, 7, 8]\n`
    );
  }

  // Test max heap
  const maxHeap = new Heap({
    priorityComp: (x: number, y: number) => x > y
  });
  for (const val of values) {
    maxHeap.push(val);
  }

  // Should pop in descending order
  const sortedValues2 = [];
  while (maxHeap.size() > 0) {
    sortedValues2.push(maxHeap.pop());
  }
  if (
    JSON.stringify(sortedValues2) !== JSON.stringify([8, 7, 6, 5, 4, 3, 2, 1])
  ) {
    throw new Error(
      `\nmax_heap popped values: got: ${sortedValues2}, want: [8, 7, 6, 5, 4, 3, 2, 1]\n`
    );
  }

  // Test heapify
  const heap = new Heap({
    priorityComp: (x: number, y: number) => x < y,
    heap: [4, 8, 2, 6, 1, 7, 3, 5],
  });
  if (heap.pop() !== 1) {
    throw new Error(`\nheap.pop(): want: 1\n`);
  }
  if (heap.pop() !== 2) {
    throw new Error(`\nheap.pop(): want: 2\n`);
  }
  if (heap.pop() !== 3) {
    throw new Error(`\nheap.pop(): want: 3\n`);
  }
}

runTests();

03-findKthLargest.ts

// https://leetcode.com/problems/kth-largest-element-in-an-array/

function findKthLargest(nums: number[], k: number): number {
  const heap = new Heap<number>({ priorityComp: (a: number, b: number) => a < b });

  for (const num of nums) {
    heap.push(num);

    if (heap.size() > k) {
      heap.pop();
    }
  }

  return heap.top();
};

04-kClosestPointToOrigin.ts

// https://leetcode.com/problems/k-closest-points-to-origin/

type Point = [number, number];
function kClosest(points: Point[], k: number): number[][] {
  const priorityComp = (a: Point, b: Point) => {
    const ad = a[0] * a[0] + a[1] * a[1];
    const bd = b[0] * b[0] + b[1] * b[1];

    const diff = ad - bd;

    // ad - bd is positive, i.e., ad has larger distance
    // from the origin than bd, thus it must be prioritized.
    if (diff > 0) return true;

    return false;
  };

  const maxHeap = new Heap<Point>({
    priorityComp,
    heap: points,
  });

  while (maxHeap.size() > k) maxHeap.pop();

  const res = [];
  while (maxHeap.size() > 0) {
    res.push(maxHeap.pop());
  }

  return res;
};