Lesson 16 of 2420 min read

Heaps and Priority Queues

JavaScript has no built-in priority queue, so you will be asked to write one. It is about forty lines, and it makes every top-k and merge-k question straightforward.

By the end: Write a heap from scratch and use it for top-k problems.

A heap answers one question fast: what is the smallest (or largest) thing here? It does not keep everything sorted — that would be more work than needed. It only guarantees that the extreme value is at the top.

Python has heapq, Java has PriorityQueue, C++ has priority_queue. JavaScript has nothing. So if you interview in JavaScript you will be asked to write one, and it is about forty lines.

A tree that lives in an array

A heap is a complete binary tree: every level is full except the last, which fills left to right. That regularity means you do not need node objects at all. Store it in an array and compute the relationships with arithmetic.

Index arithmetic replaces pointers

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24const parent = (i) => Math.floor((i - 1) / 2);
const left = (i) => 2 * i + 1;
const right = (i) => 2 * i + 2;

//        1          index 0
//      /   \
//     3     5       indices 1, 2
//    / \   /
//   4   8 7         indices 3, 4, 5
const heap = [1, 3, 5, 4, 8, 7];

console.log('array:', heap);
for (let i = 0; i < heap.length; i++) {
  const l = left(i), r = right(i);
  console.log(
    `index ${i} (value ${heap[i]})`,
    '| parent:', i === 0 ? '-' : `${heap[parent(i)]}`,
    '| children:', l < heap.length ? heap[l] : '-', r < heap.length ? heap[r] : '-'
  );
}
console.log('every parent <= its children:', heap.every((v, i) => {
  const l = left(i), r = right(i);
  return (l >= heap.length || v <= heap[l]) && (r >= heap.length || v <= heap[r]);
}));

Output

The two operations

Sift up after inserting: put the new value at the end, then swap it upward while it is smaller than its parent.

Sift down after removing the root: move the last element to the top, then swap it downward with its smaller child until the heap property holds.

Both walk one path from root to leaf, so both are O(log n).

Sift up, traced

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16function siftUpTraced(heap, i) {
  while (i > 0) {
    const p = Math.floor((i - 1) / 2);
    if (heap[p] <= heap[i]) break;                  // parent is fine, stop
    console.log(`  swap ${heap[i]} (index ${i}) with parent ${heap[p]} (index ${p})`);
    [heap[p], heap[i]] = [heap[i], heap[p]];
    i = p;
  }
}

const heap = [1, 3, 5, 4, 8, 7];
console.log('before:', [...heap]);
heap.push(0);                                        // insert 0 at the end
console.log('pushed 0:', [...heap]);
siftUpTraced(heap, heap.length - 1);
console.log('after:', heap, '← 0 rose to the root');

Output

Sift down, traced

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21function siftDownTraced(heap, i) {
  const n = heap.length;
  while (true) {
    const l = 2 * i + 1, r = 2 * i + 2;
    let smallest = i;
    if (l < n && heap[l] < heap[smallest]) smallest = l;
    if (r < n && heap[r] < heap[smallest]) smallest = r;
    if (smallest === i) break;                       // in the right place
    console.log(`  swap ${heap[i]} (index ${i}) with child ${heap[smallest]} (index ${smallest})`);
    [heap[i], heap[smallest]] = [heap[smallest], heap[i]];
    i = smallest;
  }
}

const heap = [0, 1, 5, 4, 8, 7, 3];
console.log('before:', [...heap]);
const min = heap[0];
heap[0] = heap.pop();                                // move the last to the root
console.log(`removed ${min}, moved last to root:`, [...heap]);
siftDownTraced(heap, 0);
console.log('after:', heap);

Output

The heap you should be able to write from memory

A complete MinHeap with a custom comparator

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71class Heap {
  #items = [];
  #compare;

  /* Default is a min-heap. Pass (a, b) => b - a for a max-heap, or any
     comparator for objects. This is why one class covers every use. */
  constructor(compare = (a, b) => a - b) {
    this.#compare = compare;
  }

  get size() { return this.#items.length; }
  get isEmpty() { return this.#items.length === 0; }
  peek() { return this.#items[0]; }

  push(value) {
    this.#items.push(value);
    this.#siftUp(this.#items.length - 1);
    return this;
  }

  pop() {
    if (this.#items.length === 0) return undefined;
    const top = this.#items[0];
    const last = this.#items.pop();
    if (this.#items.length > 0) {
      this.#items[0] = last;
      this.#siftDown(0);
    }
    return top;
  }

  #siftUp(i) {
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (this.#compare(this.#items[p], this.#items[i]) <= 0) break;
      [this.#items[p], this.#items[i]] = [this.#items[i], this.#items[p]];
      i = p;
    }
  }

  #siftDown(i) {
    const n = this.#items.length;
    while (true) {
      const l = 2 * i + 1, r = 2 * i + 2;
      let best = i;
      if (l < n && this.#compare(this.#items[l], this.#items[best]) < 0) best = l;
      if (r < n && this.#compare(this.#items[r], this.#items[best]) < 0) best = r;
      if (best === i) break;
      [this.#items[i], this.#items[best]] = [this.#items[best], this.#items[i]];
      i = best;
    }
  }
}

const minHeap = new Heap();
for (const v of [5, 3, 8, 1, 9, 2]) minHeap.push(v);
console.log('min-heap peek:', minHeap.peek());
const drained = [];
while (!minHeap.isEmpty) drained.push(minHeap.pop());
console.log('popped in order:', drained);

const maxHeap = new Heap((a, b) => b - a);
for (const v of [5, 3, 8, 1, 9, 2]) maxHeap.push(v);
console.log('max-heap peek:', maxHeap.peek());

// Objects, ordered by a field.
const tasks = new Heap((a, b) => a.priority - b.priority);
tasks.push({ name: 'deploy', priority: 3 });
tasks.push({ name: 'fix outage', priority: 1 });
tasks.push({ name: 'write docs', priority: 5 });
console.log('most urgent:', tasks.pop(), '| next:', tasks.pop());

Output

Costs

peek — O(1)
push — O(log n)
pop — O(log n)
Building from an array — O(n) if you heapify in place, not O(n log n)
Finding an arbitrary value — O(n). A heap is not searchable; only the top is special.
Space — O(n), and O(1) extra if you heapify in place

Heapify in O(n), which surprises people

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29// Sifting down from the last parent backwards builds a heap in linear time.
// Most nodes are near the bottom and barely move, which is why the total is O(n)
// rather than O(n log n).
function heapify(arr) {
  const n = arr.length;
  let swaps = 0;
  const siftDown = (i) => {
    while (true) {
      const l = 2 * i + 1, r = 2 * i + 2;
      let best = i;
      if (l < n && arr[l] < arr[best]) best = l;
      if (r < n && arr[r] < arr[best]) best = r;
      if (best === i) break;
      [arr[i], arr[best]] = [arr[best], arr[i]];
      swaps++;
      i = best;
    }
  };
  for (let i = Math.floor(n / 2) - 1; i >= 0; i--) siftDown(i);
  return { arr, swaps };
}

console.log(heapify([9, 4, 7, 1, 8, 2, 6, 3, 5]));

// Compare the swap count with pushing one at a time.
const data = Array.from({ length: 1000 }, (_, i) => 1000 - i);   // worst case: reversed
const built = heapify([...data]);
console.log('heapify 1000 reversed items:', built.swaps, 'swaps');
console.log('pushing one at a time would be about n log n =', Math.round(1000 * Math.log2(1000)));

Output

The reason heaps come up: top-k

The single most common heap question. To keep the k largest values, use a min-heap of size k. The smallest of your winners sits at the top, ready to be evicted the moment something better arrives.

Top k with a size-k min-heap: O(n log k)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30class Heap {
  #items = []; #cmp;
  constructor(cmp = (a, b) => a - b) { this.#cmp = cmp; }
  get size() { return this.#items.length; }
  peek() { return this.#items[0]; }
  push(v) { this.#items.push(v); let i = this.#items.length - 1;
    while (i > 0) { const p = (i - 1) >> 1; if (this.#cmp(this.#items[p], this.#items[i]) <= 0) break;
      [this.#items[p], this.#items[i]] = [this.#items[i], this.#items[p]]; i = p; } return this; }
  pop() { if (!this.#items.length) return undefined; const top = this.#items[0]; const last = this.#items.pop();
    if (this.#items.length) { this.#items[0] = last; let i = 0; const n = this.#items.length;
      while (true) { const l = 2*i+1, r = 2*i+2; let b = i;
        if (l < n && this.#cmp(this.#items[l], this.#items[b]) < 0) b = l;
        if (r < n && this.#cmp(this.#items[r], this.#items[b]) < 0) b = r;
        if (b === i) break; [this.#items[i], this.#items[b]] = [this.#items[b], this.#items[i]]; i = b; } }
    return top; }
  toArray() { return [...this.#items]; }
}

function topK(nums, k) {
  const heap = new Heap();                 // MIN-heap, holding the k largest
  for (const x of nums) {
    heap.push(x);
    if (heap.size > k) heap.pop();         // evict the smallest winner
  }
  return heap.toArray().sort((a, b) => b - a);
}

console.log('top 3 of [3,2,1,5,6,4]:', topK([3, 2, 1, 5, 6, 4], 3));
console.log('top 2 of [3,2,3,1,2,4,5,5,6]:', topK([3,2,3,1,2,4,5,5,6], 2));
console.log('k = n:', topK([1, 2, 3], 3));

Output

K most frequent elements: count, then heap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33class Heap {
  #a = []; #c;
  constructor(c = (x, y) => x - y) { this.#c = c; }
  get size() { return this.#a.length; }
  push(v) { this.#a.push(v); let i = this.#a.length - 1;
    while (i > 0) { const p = (i-1)>>1; if (this.#c(this.#a[p], this.#a[i]) <= 0) break;
      [this.#a[p], this.#a[i]] = [this.#a[i], this.#a[p]]; i = p; } return this; }
  pop() { if (!this.#a.length) return undefined; const t = this.#a[0], l = this.#a.pop();
    if (this.#a.length) { this.#a[0] = l; let i = 0, n = this.#a.length;
      while (true) { const L = 2*i+1, R = 2*i+2; let b = i;
        if (L < n && this.#c(this.#a[L], this.#a[b]) < 0) b = L;
        if (R < n && this.#c(this.#a[R], this.#a[b]) < 0) b = R;
        if (b === i) break; [this.#a[i], this.#a[b]] = [this.#a[b], this.#a[i]]; i = b; } }
    return t; }
  toArray() { return [...this.#a]; }
}

function topKFrequent(nums, k) {
  const counts = new Map();
  for (const x of nums) counts.set(x, (counts.get(x) ?? 0) + 1);

  // Min-heap on the count, holding the k most frequent.
  const heap = new Heap((a, b) => a[1] - b[1]);
  for (const entry of counts) {
    heap.push(entry);
    if (heap.size > k) heap.pop();
  }
  return heap.toArray().sort((a, b) => b[1] - a[1]).map(([value]) => value);
}

console.log(topKFrequent([1, 1, 1, 2, 2, 3], 2));    // [1,2]
console.log(topKFrequent([1], 1));                    // [1]
console.log(topKFrequent([4,4,4,5,5,6], 2));          // [4,5]

Output

Merging k sorted sequences

K-way merge with a heap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39class Heap {
  #a = []; #c;
  constructor(c) { this.#c = c; }
  get size() { return this.#a.length; }
  push(v) { this.#a.push(v); let i = this.#a.length - 1;
    while (i > 0) { const p = (i-1)>>1; if (this.#c(this.#a[p], this.#a[i]) <= 0) break;
      [this.#a[p], this.#a[i]] = [this.#a[i], this.#a[p]]; i = p; } return this; }
  pop() { const t = this.#a[0], l = this.#a.pop();
    if (this.#a.length) { this.#a[0] = l; let i = 0, n = this.#a.length;
      while (true) { const L = 2*i+1, R = 2*i+2; let b = i;
        if (L < n && this.#c(this.#a[L], this.#a[b]) < 0) b = L;
        if (R < n && this.#c(this.#a[R], this.#a[b]) < 0) b = R;
        if (b === i) break; [this.#a[i], this.#a[b]] = [this.#a[b], this.#a[i]]; i = b; } }
    return t; }
}

// Keep one candidate from each list. Pop the smallest, then refill from the
// list it came from. O(N log k) where N is the total number of items.
function mergeKSorted(lists) {
  const heap = new Heap((a, b) => a.value - b.value);
  lists.forEach((list, listIndex) => {
    if (list.length > 0) heap.push({ value: list[0], listIndex, position: 0 });
  });

  const out = [];
  while (heap.size > 0) {
    const { value, listIndex, position } = heap.pop();
    out.push(value);
    const next = position + 1;
    if (next < lists[listIndex].length) {
      heap.push({ value: lists[listIndex][next], listIndex, position: next });
    }
  }
  return out;
}

console.log(mergeKSorted([[1, 4, 5], [1, 3, 4], [2, 6]]));
console.log(mergeKSorted([[], [1], []]));
console.log(mergeKSorted([]));

Output

Two heaps: the running median

An elegant trick. Keep a max-heap of the lower half and a min-heap of the upper half, balanced in size. The median is then at one or both tops.

Median of a stream in O(log n) per insert

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44class Heap {
  #a = []; #c;
  constructor(c) { this.#c = c; }
  get size() { return this.#a.length; }
  peek() { return this.#a[0]; }
  push(v) { this.#a.push(v); let i = this.#a.length - 1;
    while (i > 0) { const p = (i-1)>>1; if (this.#c(this.#a[p], this.#a[i]) <= 0) break;
      [this.#a[p], this.#a[i]] = [this.#a[i], this.#a[p]]; i = p; } return this; }
  pop() { const t = this.#a[0], l = this.#a.pop();
    if (this.#a.length) { this.#a[0] = l; let i = 0, n = this.#a.length;
      while (true) { const L = 2*i+1, R = 2*i+2; let b = i;
        if (L < n && this.#c(this.#a[L], this.#a[b]) < 0) b = L;
        if (R < n && this.#c(this.#a[R], this.#a[b]) < 0) b = R;
        if (b === i) break; [this.#a[i], this.#a[b]] = [this.#a[b], this.#a[i]]; i = b; } }
    return t; }
}

class MedianFinder {
  #low = new Heap((a, b) => b - a);    // max-heap: largest of the small half
  #high = new Heap((a, b) => a - b);   // min-heap: smallest of the big half

  add(num) {
    // Always push to low first, then move its top across. Keeps the invariant
    // that everything in low is <= everything in high.
    this.#low.push(num);
    this.#high.push(this.#low.pop());

    // Rebalance so low has the same size, or one more.
    if (this.#high.size > this.#low.size) this.#low.push(this.#high.pop());
    return this;
  }

  median() {
    if (this.#low.size === 0) return null;
    if (this.#low.size > this.#high.size) return this.#low.peek();
    return (this.#low.peek() + this.#high.peek()) / 2;
  }
}

const m = new MedianFinder();
for (const x of [5, 15, 1, 3]) {
  m.add(x);
  console.log(`after adding ${String(x).padStart(2)}: median is ${m.median()}`);
}

Output

Heap sort

Sorting in place with O(1) extra space

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30function heapSort(arr) {
  const n = arr.length;

  const siftDown = (i, limit) => {
    while (true) {
      const l = 2 * i + 1, r = 2 * i + 2;
      let biggest = i;
      if (l < limit && arr[l] > arr[biggest]) biggest = l;
      if (r < limit && arr[r] > arr[biggest]) biggest = r;
      if (biggest === i) break;
      [arr[i], arr[biggest]] = [arr[biggest], arr[i]];
      i = biggest;
    }
  };

  // 1. Build a MAX-heap in place, O(n).
  for (let i = Math.floor(n / 2) - 1; i >= 0; i--) siftDown(i, n);

  // 2. Repeatedly move the largest to the end and shrink the heap.
  for (let end = n - 1; end > 0; end--) {
    [arr[0], arr[end]] = [arr[end], arr[0]];
    siftDown(0, end);
  }
  return arr;
}

console.log(heapSort([5, 2, 9, 1, 6, 3]));
console.log(heapSort([1]));
console.log(heapSort([3, 3, 3]));
console.log(heapSort([9, 8, 7, 6, 5]));

Output

What a heap cannot do

Search. Finding an arbitrary value is O(n). Use a Map if you need that.
Give you sorted order cheaply. You must pop everything, which is O(n log n).
Update a value in place. Not without an index Map alongside it, which is what a real Dijkstra implementation does.
Tell you the maximum, in a min-heap. The bottom is unordered; the maximum could be any leaf.

Your turn: four heap problems

The Heap class is provided so you can focus on the algorithm.

kthLargest(nums, k) — using a size-k min-heap.
lastStoneWeight(stones) — repeatedly smash the two heaviest together; the survivor is their difference. Return what is left.
kClosestToOrigin(points, k) — the k points closest to (0,0). Compare squared distance and skip the square root.
minMeetingRooms(intervals) — how many rooms are needed so no two meetings overlap. Sort by start, and keep a min-heap of end times.

Editable. Fill them in, then run.

class Heap {
  #a = []; #c;
  constructor(c = (x, y) => x - y) { this.#c = c; }
  get size() { return this.#a.length; }
  get isEmpty() { return this.#a.length === 0; }
  peek() { return this.#a[0]; }
  push(v) { this.#a.push(v); let i = this.#a.length - 1;
    while (i > 0) { const p = (i-1)>>1; if (this.#c(this.#a[p], this.#a[i]) <= 0) break;
      [this.#a[p], this.#a[i]] = [this.#a[i], this.#a[p]]; i = p; } return this; }
  pop() { if (!this.#a.length) return undefined; const t = this.#a[0], l = this.#a.pop();
    if (this.#a.length) { this.#a[0] = l; let i = 0; const n = this.#a.length;
      while (true) { const L = 2*i+1, R = 2*i+2; let b = i;
        if (L < n && this.#c(this.#a[L], this.#a[b]) < 0) b = L;
        if (R < n && this.#c(this.#a[R], this.#a[b]) < 0) b = R;
        if (b === i) break; [this.#a[i], this.#a[b]] = [this.#a[b], this.#a[i]]; i = b; } }
    return t; }
  toArray() { return [...this.#a]; }
}

function kthLargest(nums, k) {
  // size-k min-heap
}

function lastStoneWeight(stones) {
  // max-heap; pop two, push the difference if non-zero
}

function kClosestToOrigin(points, k) {
  // max-heap of size k on squared distance
}

function minMeetingRooms(intervals) {
  // sort by start; min-heap of end times; pop when a room frees up
}

console.log(kthLargest([3, 2, 1, 5, 6, 4], 2));                  // 5
console.log(kthLargest([3,2,3,1,2,4,5,5,6], 4));                 // 4
console.log(lastStoneWeight([2, 7, 4, 1, 8, 1]));                // 1
console.log(lastStoneWeight([1]));                                // 1
console.log(kClosestToOrigin([[1,3],[-2,2]], 1));                // [[-2,2]]
console.log(kClosestToOrigin([[3,3],[5,-1],[-2,4]], 2));         // [[3,3],[-2,4]] in some order
console.log(minMeetingRooms([[0,30],[5,10],[15,20]]));           // 2
console.log(minMeetingRooms([[7,10],[2,4]]));                    // 1

Output

Recap

A heap guarantees only that the extreme value is on top; it is not sorted and it cannot be searched.
It is a complete tree stored in an array: parent is (i-1)>>1, children are 2i+1 and 2i+2.
push and pop are O(log n), peek is O(1), and heapifying an existing array is O(n).
Write one class taking a comparator so it serves as both min-heap and max-heap.
For the k largest use a size-k min-heap so you can evict the weakest winner; two heaps facing each other give a running median.

A tree that lives in an array

The two operations

The heap you should be able to write from memory

Costs

The reason heaps come up: top-k

Merging k sorted sequences

Two heaps: the running median

Heap sort

What a heap cannot do

Your turn: four heap problems

Recap

New writing, when it lands.