Heaps: Top-K, K-way Merge and Two Heaps

// Template 1: a binary heap. JavaScript has no built-in priority queue, so write this once and reuse it.
// The heap lives in an array: the parent of index i is (i - 1) >> 1, its children are 2i + 1 and 2i + 2.
// compare(a, b) < 0 means a leaves first: (a, b) => a - b is a min-heap, (a, b) => b - a a max-heap.
export class Heap {
  constructor(compare = (a, b) => a - b) {
    this.items = [];
    this.compare = compare;
  }

  size() {
    return this.items.length;
  }

  peek() {
    return this.items[0]; // undefined when empty
  }

  push(value) {
    this.items.push(value);              // add as the last leaf…
    this.siftUp(this.items.length - 1);  // …and float it up past larger parents
  }

  pop() {
    const items = this.items;
    if (items.length === 0) return undefined;
    const top = items[0];
    const last = items.pop();
    if (items.length > 0) {
      items[0] = last;                   // move the last leaf to the root…
      this.siftDown(0);                  // …and sink it below smaller children
    }
    return top;
  }

  siftUp(i) {
    const items = this.items;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(items[i], items[parent]) >= 0) break; // parent already comes first
      [items[i], items[parent]] = [items[parent], items[i]];
      i = parent;
    }
  }

  siftDown(i) {
    const items = this.items;
    for (;;) {
      const left = 2 * i + 1;
      const right = left + 1;
      let best = i;                      // whichever of i, left, right should come first
      if (left < items.length && this.compare(items[left], items[best]) < 0) best = left;
      if (right < items.length && this.compare(items[right], items[best]) < 0) best = right;
      if (best === i) return;
      [items[i], items[best]] = [items[best], items[i]];
      i = best;
    }
  }
}

import { Heap } from './heap.js';

// Template 2: top k. To keep the k LARGEST values, keep a MIN-heap of size k: its top is the
// weakest of the current top k, the one to evict when something bigger arrives. O(n log k).
export function kLargest(nums, k) {
  const heap = new Heap((a, b) => a - b);
  for (const x of nums) {
    heap.push(x);
    if (heap.size() > k) heap.pop();   // drop the smallest: it is not in the top k
  }
  const out = [];
  while (heap.size() > 0) out.push(heap.pop());
  return out.reverse();                // largest first
}

// The mirror image: the k SMALLEST items by some key need a MAX-heap of size k.
export function kSmallestBy(items, k, key) {
  const heap = new Heap((a, b) => key(b) - key(a));
  for (const item of items) {
    heap.push(item);
    if (heap.size() > k) heap.pop();   // drop the largest key
  }
  const out = [];
  while (heap.size() > 0) out.push(heap.pop());
  return out.reverse();                // smallest key first
}

import { Heap } from './heap.js';

// Template 3: k-way merge. The next output is always the smallest of the k current heads,
// so keep one entry per input in a min-heap. Each of the N items is pushed and popped once: O(N log k).
export function mergeSortedArrays(arrays) {
  const heap = new Heap((a, b) => a.value - b.value);
  arrays.forEach((arr, list) => {
    if (arr.length > 0) heap.push({ value: arr[0], list, index: 0 });
  });
  const out = [];
  while (heap.size() > 0) {
    const { value, list, index } = heap.pop();
    out.push(value);
    const next = index + 1;
    if (next < arrays[list].length) heap.push({ value: arrays[list][next], list, index: next }); // advance that input
  }
  return out;
}

import { Heap } from './heap.js';

// Template 4: two heaps. `low` (a max-heap) holds the smaller half, `high` (a min-heap) the larger half,
// with low.size() === high.size() or one more. The median is then low's top, or the average of both tops.
export function runningMedians(nums) {
  const low = new Heap((a, b) => b - a);
  const high = new Heap((a, b) => a - b);
  const medians = [];
  for (const x of nums) {
    low.push(x);
    high.push(low.pop());                                       // keeps every low value <= every high value
    if (high.size() > low.size()) low.push(high.pop());         // keeps the sizes balanced
    medians.push(low.size() > high.size() ? low.peek() : (low.peek() + high.peek()) / 2);
  }
  return medians;
}