DSA for Interviews, From Zero

Part 12 · Lesson 46 · read

Heaps and Top-K

Build a priority queue by hand in JavaScript, then use it to find the top K items.

On this page

  1. The problem: always give me the smallest
  2. The heap rule and its array form
  3. Push and pop: sift up, sift down
  4. Traced: six pushes and one pop
  5. Writing a Heap in JavaScript
  6. Heapify: building a heap in O(n)
  7. Top-K: the size-k heap
  8. Top k frequent elements
  9. Merge k sorted lists
  10. Median of a data stream: two heaps
  11. Practice questions (7)
  12. Make it stick
  13. What's next

The problem: always give me the smallest

Some jobs keep asking for "the smallest item so far" (or the largest) while new items keep arriving. Examples: the next task with the earliest deadline, the nearest of many points, or the lightest stone. A tool that can add an item and remove the best item is called a priority queue (a waiting line where the most important item goes first, not the oldest). How could you build one from what you know?

  • An unsorted array: adding is O(1), but finding the minimum means looking at every item, which is O(n).
  • A sorted array: the minimum is at the end, but adding in the right place moves many items, which is O(n).

Either way, one of the two jobs is slow (O(n) means the work grows with the number of items). A heap makes both jobs O(log n), which is much faster. Looking at the best item is O(1), which means it takes the same small number of steps, however many items there are. A heap is the standard way to build a priority queue.

JavaScript has no built-in heap
Python has heapq and Java has PriorityQueue, but standard JavaScript has neither. In an interview you write your own, so this lesson builds one from scratch. It is only about 40 lines. You will paste the same class into every solution in this lesson.

The heap property and its array form

A binary heap is a binary tree (each item has at most two children below it) with two rules:

  • Shape: the tree is complete. Every level is full, except maybe the last one. The last level is filled from the left. There are no gaps.
  • Heap property: in a min-heap, every item is smaller than or equal to its children. A max-heap is the opposite: every item is bigger than or equal to its children. So in a min-heap, the smallest item is always at the top (the root).

Notice what the rule does not say. It does not say how items on the same level compare with each other. For example, in a min-heap the left child may be bigger or smaller than the right child. A heap is only partly sorted. It is sorted just enough to know the minimum. That is why it is cheaper to keep than a fully sorted list.

Because the tree has no gaps, it fits in a plain array. You store it level by level, and you need no pointers. For an item at index i:

Dry run · the index formulas, shown on the min-heap [1, 3, 2, 5, 9, 8]
Index iValueParent (i − 1) >> 1Children 2i + 1, 2i + 2
01none (root)1 and 2 (values 3, 2)
130 (value 1)3 and 4 (values 5, 9)
220 (value 1)5 and 6 (value 8, and 6 does not exist)
351 (value 3)7 and 8 (neither exists)
491 (value 3)9 and 10 (neither exists)
582 (value 2)11 and 12 (neither exists)

The expression (i − 1) >> 1 means 'divide by two and drop the decimal part'. (It shifts the bits of the number one place to the right.) So (5 − 1) >> 1 = 2 and (4 − 1) >> 1 = 1. Math.floor((i − 1) / 2) means the same thing.

A complete tree with n items has a height of about log₂ n (the number of times you can halve n). Every operation below walks along one path from the root down to a leaf (an item with no children). That is where O(log n) comes from. With a million items, that is only about 20 steps.

Push and pop: sift up, sift down

Both operations break the heap rule in one place. Then they fix it by walking along a single path.

Push (sift up). Put the new item in the first free slot at the end of the array. This keeps the shape complete. The new item may now be smaller than its parent. So compare it with its parent, and swap while the parent is bigger. This moving up is called sift up (or "bubble up", like a bubble rising in water). It stops at the root, or at the first parent that is not bigger.

Pop (sift down). The answer is at index 0. But taking it out leaves a hole at the root. Fill the hole with the last item of the array. This keeps the shape complete, and the array just gets one item shorter. That item is probably too big for the root. So compare it with its two children. Swap it with the smaller child, and repeat while a child is smaller. This is sift down. Using the smaller child matters. That child becomes the new parent of the other child, so the rule still holds between them.

Each walk is at most as long as the height of the tree. So push and pop are both O(log n). Looking at index 0 is O(1).

Traced: six pushes and one pop

step 1 / 24
Set up
heap = []

A min-heap kept in an array. The item at index i has its parent at (i − 1) >> 1 and its children at 2i + 1 and 2i + 2.

▶const heap = [];2for (const x of [5, 3, 8, 1, 9, 2]) {3  heap.push(x);4  let i = heap.length - 1;5  while (i > 0) {6    const parent = (i - 1) >> 1;7    if (heap[parent] <= heap[i]) break;8    [heap[parent], heap[i]] = [heap[i], heap[parent]];9    i = parent;10  }11}12const top = heap[0];13const last = heap.pop();14heap[0] = last;15let j = 0;16while (true) {17  const l = 2 * j + 1, r = 2 * j + 2;18  let small = j;19  if (l < heap.length && heap[l] < heap[small]) small = l;20  if (r < heap.length && heap[r] < heap[small]) small = r;21  if (small === j) break;22  [heap[small], heap[j]] = [heap[j], heap[small]];23  j = small;24}
Variables
heap[]
Console
(nothing printed yet)
Pushing 1 sifts all the way up to the root, because 1 is the new minimum. The pop then moves the last item (8) to the root and sifts it down just one level. · interactive

Writing a Heap in JavaScript

The class below takes a comparator. This is a small function, compare(a, b), that tells the heap which of two items should come out first. It returns a negative number when a should leave the heap before b. The default, a − b, gives a min-heap of numbers. Pass (a, b) => b − a and the very same code becomes a max-heap. Pass (a, b) => a[1] − b[1] and it orders pairs by their second value. So you write one class and you never copy the logic.

JavaScript
class Heap {  // compare(a, b) < 0 means "a comes out before b".  // The default makes a min-heap of numbers. Pass (a, b) => b - a for a max-heap.  constructor(compare = (a, b) => a - b) {    this.data = [];    this.compare = compare;  }  get size() { return this.data.length; }  peek() { return this.data[0]; }                     // the top item, without removing it  push(value) {    const d = this.data;    d.push(value);                                    // put it in the first free slot...    let i = d.length - 1;    while (i > 0) {                                   // ...and sift it up      const parent = (i - 1) >> 1;      if (this.compare(d[i], d[parent]) >= 0) break;  // parent is fine: stop      [d[i], d[parent]] = [d[parent], d[i]];      i = parent;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0];    const last = d.pop();                             // take the last item off the end...    if (d.length > 0) {      d[0] = last;                                    // ...put it at the root, and sift it down      let i = 0;      while (true) {        const left = 2 * i + 1, right = 2 * i + 2;        let best = i;        if (left < d.length && this.compare(d[left], d[best]) < 0) best = left;        if (right < d.length && this.compare(d[right], d[best]) < 0) best = right;        if (best === i) break;                        // both children are no better: stop        [d[i], d[best]] = [d[best], d[i]];        i = best;      }    }    return top;  }}const minHeap = new Heap();for (const x of [5, 3, 8, 1, 9, 2]) minHeap.push(x);console.log(minHeap.peek());                          // 1const sorted = [];while (minHeap.size > 0) sorted.push(minHeap.pop());console.log(sorted);                                  // [1, 2, 3, 5, 8, 9]   (n pops = heap sort, O(n log n))const maxHeap = new Heap((a, b) => b - a);for (const x of [5, 3, 8, 1, 9, 2]) maxHeap.push(x);console.log(maxHeap.pop(), maxHeap.pop());            // 9 8const byCount = new Heap((a, b) => a[1] - b[1]);      // items can be anything, if compare says how to order thembyCount.push(["a", 4]); byCount.push(["b", 1]); byCount.push(["c", 3]);console.log(byCount.pop());                           // [ 'b', 1 ]

If you take all n items out with pop, they come out in sorted order. This is called heap sort, and it is O(n log n). You rarely write it. But it shows that the heap really does ordering work.

Do not read the heap array as sorted
heap.data only promises that the smallest item is first. If you print it, loop over it, or use data[1] as the second smallest, you have a bug. Always take items out with pop.

Heapify: building a heap in O(n)

Say you already have all n items. Pushing them one by one costs O(n log n). There is a faster way, called heapify. Treat the array as a tree, and sift down every item that has children. Start from the last one and move back to the root. The leaves (the second half of the array) are already valid one-item heaps. So the loop starts at index n/2 − 1.

JavaScript
function siftDown(a, i, n) {                  // push a[i] down within the first n items  while (true) {    const left = 2 * i + 1, right = 2 * i + 2;    let small = i;    if (left < n && a[left] < a[small]) small = left;    if (right < n && a[right] < a[small]) small = right;    if (small === i) return;    [a[i], a[small]] = [a[small], a[i]];    i = small;  }}function heapify(a) {                         // rearranges the array in place into a min-heap  for (let i = (a.length >> 1) - 1; i >= 0; i--) siftDown(a, i, a.length);  return a;                                   // items from index n/2 onwards are leaves already}console.log(heapify([5, 3, 8, 1, 9, 2])); // [1, 3, 2, 5, 9, 8]console.log(heapify([9, 8, 7, 6, 5]));    // [5, 6, 7, 9, 8]

Why is this O(n) and not O(n log n)? A sift down costs at most the height of the item it starts from. Most items are near the bottom. About n/2 items are leaves (cost 0). About n/4 are one level up (cost at most 1). About n/8 cost at most 2, and so on. The sum n/4 · 1 + n/8 · 2 + n/16 · 3 + … stays below n. Remember this idea: few items are tall, and many are short.

Top-K: the size-k heap

Find the k-th largest item. Sorting everything costs O(n log n). Instead, keep a heap that never grows past k items. Which kind of heap? A min-heap, which is surprising. The heap holds the k largest items seen so far. Its top is the smallest of those k. That is the first item to throw out when a bigger one arrives. After you have seen every item, the top is the k-th largest overall.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function kthLargest(nums, k) {  const heap = new Heap();                    // a MIN-heap that never holds more than k items  for (const x of nums) {    heap.push(x);    if (heap.size > k) heap.pop();            // evict the smallest: it cannot be among the k largest  }  return heap.peek();                         // the smallest of the k largest = the k-th largest}console.log(kthLargest([3, 2, 1, 5, 6, 4], 2));          // 5console.log(kthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4

Each of the n items costs O(log k) to push and maybe pop. So the total is O(n log k) time and O(k) space. When k is small, this is much better than sorting. It also works on a stream (items that keep arriving one at a time, so you cannot store everything). Remember the rule: for the k largest, use a min-heap of size k. For the k smallest, use a max-heap of size k.

Top k frequent elements

First count how often each value appears, using a Map (a lookup table that stores key and value pairs, here value and count). Then use the same size-k trick on the counts. The heap holds [value, count] pairs, and the comparator looks at the count. So the least frequent of the current candidates is always on top, ready to be thrown out.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function topKFrequent(nums, k) {  const count = new Map();  for (const x of nums) count.set(x, (count.get(x) ?? 0) + 1);  const heap = new Heap((a, b) => a[1] - b[1]);         // min-heap of [value, count], ordered by count  for (const entry of count) {    heap.push(entry);    if (heap.size > k) heap.pop();                      // drop the least frequent of the k + 1  }  const result = [];  while (heap.size > 0) result.push(heap.pop()[0]);  return result.reverse();                              // most frequent first}console.log(topKFrequent([1, 1, 1, 2, 2, 3], 2)); // [1, 2]console.log(topKFrequent([4, 4, 5, 5, 5, 6], 1)); // [5]

Counting is O(n). Pushing the different values is O(d log k), where d is the number of different values. So the total is O(n + d log k), which is at most O(n log k). (A bucket-sort version can reach O(n). You will see it in the practice questions.)

Merge k sorted lists

Lesson 37 merged two sorted linked lists (chains of nodes, where each node points to the next one) by repeatedly taking the smaller front node. With k lists, the next node of the answer is the smallest of k front nodes. Comparing all k fronts each time costs O(k) per node. A min-heap of the current fronts finds the smallest in O(log k). Pop the smallest, attach it to the answer, and push the next node from the same list.

JavaScript
class ListNode {  constructor(val, next = null) { this.val = val; this.next = next; }}function fromArray(values) {  const dummy = new ListNode(0);  let tail = dummy;  for (const v of values) { tail.next = new ListNode(v); tail = tail.next; }  return dummy.next;}function toArray(head) {  const out = [];  for (let cur = head; cur !== null; cur = cur.next) out.push(cur.val);  return out;}class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function mergeKLists(lists) {  const heap = new Heap((a, b) => a.val - b.val);   // holds the current front node of each list  for (const head of lists) if (head !== null) heap.push(head);  const dummy = new ListNode(0);  let tail = dummy;  while (heap.size > 0) {    const node = heap.pop();                        // the smallest front node of all lists    tail.next = node;    tail = node;    if (node.next !== null) heap.push(node.next);   // its list has a new front  }  return dummy.next;}const lists = [fromArray([1, 4, 5]), fromArray([1, 3, 4]), fromArray([2, 6])];console.log(toArray(mergeKLists(lists)));           // [1, 1, 2, 3, 4, 4, 5, 6]console.log(toArray(mergeKLists([null, fromArray([7])]))); // [7]

The heap never holds more than k nodes, and each of the N nodes is pushed and popped once. So the time is O(N log k) and the extra space is O(k). The nodes are linked again, not copied.

Median of a data stream: two heaps

Numbers arrive one at a time. At any moment you must report the median. The median is the middle value of the sorted data. If the count is even, it is the average of the two middle values. Sorting again after each number is too slow. The trick is to split the data into a lower half and an upper half:

  • the lower half lives in a max-heap, so its biggest item is on top;
  • the upper half lives in a min-heap, so its smallest item is on top;
  • the two halves have the same size, or the lower half has one extra item.

Now the median comes from the two tops. If the count is odd, it is the top of the lower half. If the count is even, it is the average of both tops. To add a number, push it into the lower heap, then pop the top of the lower heap into the upper heap. This makes sure every lower item stays smaller than or equal to every upper item. Then, if the upper half became bigger, move one item back.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}class MedianFinder {  constructor() {    this.low = new Heap((a, b) => b - a);    // max-heap: the smaller half, biggest on top    this.high = new Heap();                  // min-heap: the larger half, smallest on top  }  addNum(x) {    this.low.push(x);    this.high.push(this.low.pop());          // route through low so every item in low stays <= every item in high    if (this.high.size > this.low.size) this.low.push(this.high.pop());   // keep low the same size or one bigger  }  findMedian() {    if (this.low.size > this.high.size) return this.low.peek();           // odd count: the middle item    return (this.low.peek() + this.high.peek()) / 2;                      // even count: average of the two middles  }}const m = new MedianFinder();const medians = [];for (const x of [5, 2, 8, 1, 9]) { m.addNum(x); medians.push(m.findMedian()); }console.log(medians); // [5, 3.5, 5, 3.5, 5]

Adding is O(log n). Reading the median is O(1).

Practice questions

When a question says "k largest", "k closest", "most frequent" or "repeatedly take the best", think of a heap.

Q1Kth largest element in an arrayMedium

Return the k-th largest item of an unsorted array. This means the k-th when sorted from biggest to smallest, not the k-th different value.

Example 1
Input: nums = [3, 2, 1, 5, 6, 4], k = 2
Output: 5
Explanation: Sorted from biggest to smallest: 6, 5, 4, 3, 2, 1. The 2nd is 5.
Example 2
Input: nums = [3, 2, 3, 1, 2, 4, 5, 5, 6], k = 4
Output: 4
Explanation: Sorted from biggest to smallest: 6, 5, 5, 4, ... Repeated numbers count separately, so the 4th is 4.
Stuck? Open one hint at a time
Hint 1
The easy way is to sort. How much work is that? Can you avoid putting everything in order?
Hint 2
You only need to remember the k biggest items. Which one of them would you throw away first?
Hint 3
A min-heap that never holds more than k items has the smallest of the k largest on top.
Show all 3 solutions with explanations
Approach 1Sort
The idea

Sort from biggest to smallest, then take the item at index k − 1.

JavaScript
function findKthLargest(nums, k) {  return [...nums].sort((a, b) => b - a)[k - 1];}console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2));          // 5console.log(findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4
How it works

It takes O(n log n) time. It is always a good first answer, and it is only two lines.

Approach 2Min-heap of size k
The idea
  1. Push each number into a min-heap.
  2. If the heap has more than k items, pop the smallest.
  3. At the end, the top is the k-th largest.
JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function findKthLargest(nums, k) {  const heap = new Heap();  for (const x of nums) {    heap.push(x);    if (heap.size > k) heap.pop();  }  return heap.peek();}console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2));          // 5console.log(findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4
How it works

It takes O(n log k) time and O(k) space. After the loop, the heap holds exactly the k largest numbers. The smallest of them is the k-th largest.

Approach 3Quickselect
The idea
  1. Pick a random number (the pivot) and split the numbers into bigger, equal and smaller groups.
  2. If the k-th place falls inside the "bigger" group, continue with only that group.
  3. If it falls in the "equal" group, the pivot is the answer.
  4. Otherwise, continue with the "smaller" group. Lower k by the sizes of the groups you skipped.
JavaScript
function findKthLargest(nums, k) {  let items = nums;  while (true) {    const pivot = items[Math.floor(Math.random() * items.length)];    const bigger = items.filter((x) => x > pivot);    const equal = items.filter((x) => x === pivot);    if (k <= bigger.length) {      items = bigger;    } else if (k <= bigger.length + equal.length) {      return pivot;    } else {      k -= bigger.length + equal.length;      items = items.filter((x) => x < pivot);    }  }}console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2));          // 5console.log(findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4
How it works

On average it takes O(n) time. Each round throws away a good part of the items, so the work keeps shrinking. The worst case is O(n²) if the pivots are unlucky. A random pivot makes this very unlikely. It uses O(n) extra space here because of the filtered arrays.

Which approach should you use?

Start with the sort to show a correct answer. Then give the size-k heap, which is the pattern interviewers expect. Mention quickselect as the option that is O(n) on average. (LeetCode 215.)

Q2Kth largest element in a streamEasy

Design a class that is created with a number k and a starting array. Its add(val) method adds a number and returns the k-th largest of all numbers so far.

Example 1
Input: KthLargest(3, [4, 5, 8, 2]), add(3), add(5), add(10), add(9), add(4)
Output: 4, 5, 5, 8, 8
Explanation: After adding 3 the numbers are 2, 3, 4, 5, 8 and the 3rd largest is 4. After adding 5 the 3rd largest is 5, and so on.
Stuck? Open one hint at a time
Hint 1
This is question 1, but the items arrive over time. What do you need to remember between calls?
Hint 2
Keep a min-heap of the k largest numbers you have seen so far.
Show all 2 solutions with explanations
Approach 1Keep a sorted array
The idea

Put each new number in its sorted place. The answer is the item k places from the end.

JavaScript
class KthLargest {  constructor(k, nums) {    this.k = k;    this.sorted = [...nums].sort((a, b) => a - b);  }  add(val) {    let i = 0;    while (i < this.sorted.length && this.sorted[i] < val) i++;    this.sorted.splice(i, 0, val);                     // shifting makes this O(n)    return this.sorted[this.sorted.length - this.k];  }}const kth = new KthLargest(3, [4, 5, 8, 2]);console.log([3, 5, 10, 9, 4].map((x) => kth.add(x)));  // [4, 5, 5, 8, 8]
How it works

Each add takes O(n) time, and the array keeps every number forever.

Approach 2Min-heap of size k
The idea

Keep only the k largest numbers in a min-heap. The top is always the k-th largest.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}class KthLargest {  constructor(k, nums) {    this.k = k;    this.heap = new Heap();    for (const x of nums) this.add(x);  }  add(val) {    this.heap.push(val);    if (this.heap.size > this.k) this.heap.pop();    return this.heap.peek();  }}const kth = new KthLargest(3, [4, 5, 8, 2]);console.log([3, 5, 10, 9, 4].map((x) => kth.add(x)));  // [4, 5, 5, 8, 8]
How it works

Each add takes O(log k) time, and the memory is O(k) in total. The problem promises there are at least k numbers when you ask. So the heap is full whenever add returns.

Which approach should you use?

Use the size-k min-heap. The stream may never end, and this stores only k numbers. (LeetCode 703.)

Q3Top k frequent elementsMedium

Return the k most frequent values of an array. The answer can be in any order. The problem promises that the answer is unique.

Example 1
Input: nums = [1, 1, 1, 2, 2, 3], k = 2
Output: [1, 2]
Explanation: 1 appears three times, 2 appears twice, and 3 appears once.
Example 2
Input: nums = [1], k = 1
Output: [1]
Explanation: There is only one value.
Stuck? Open one hint at a time
Hint 1
First, count each value with a Map.
Hint 2
Then it is a top-k problem on the counts. A min-heap that holds at most k items keeps the most frequent ones.
Hint 3
A count is never more than n. Could you put values into buckets numbered by their count?
Show all 3 solutions with explanations
Approach 1Sort by count
The idea

Count the values. Sort the different values by count, biggest first. Take the first k.

JavaScript
function topKFrequent(nums, k) {  const count = new Map();  for (const x of nums) count.set(x, (count.get(x) ?? 0) + 1);  return [...count.keys()].sort((a, b) => count.get(b) - count.get(a)).slice(0, k);}console.log(topKFrequent([1, 1, 1, 2, 2, 3], 2)); // [1, 2]console.log(topKFrequent([1], 1));                // [1]
How it works

It takes O(n + d log d) time, where d is the number of different values. That is at most O(n log n).

Approach 2Min-heap of size k
The idea

Push each [value, count] pair into a min-heap ordered by count. Pop whenever it has more than k items.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function topKFrequent(nums, k) {  const count = new Map();  for (const x of nums) count.set(x, (count.get(x) ?? 0) + 1);  const heap = new Heap((a, b) => a[1] - b[1]);  for (const entry of count) {    heap.push(entry);    if (heap.size > k) heap.pop();  }  const result = [];  while (heap.size > 0) result.push(heap.pop()[0]);  return result.reverse();}console.log(topKFrequent([1, 1, 1, 2, 2, 3], 2)); // [1, 2]console.log(topKFrequent([1], 1));                // [1]
How it works

It takes O(n + d log k) time. It is good when k is much smaller than d.

Approach 3Bucket sort by count
The idea
  1. Count each value.
  2. Make n + 1 buckets. Put each value in the bucket with the same number as its count. A value cannot appear more than n times.
  3. Read the buckets from the highest number down until you have k values.
JavaScript
function topKFrequent(nums, k) {  const count = new Map();  for (const x of nums) count.set(x, (count.get(x) ?? 0) + 1);  const buckets = Array.from({ length: nums.length + 1 }, () => []);  for (const [value, c] of count) buckets[c].push(value);  const result = [];  for (let c = buckets.length - 1; c >= 0 && result.length < k; c--) {    for (const value of buckets[c]) {      if (result.length < k) result.push(value);    }  }  return result;}console.log(topKFrequent([1, 1, 1, 2, 2, 3], 2)); // [1, 2]console.log(topKFrequent([1], 1));                // [1]
How it works

It takes O(n) time and space, with no heap and no sort. It works because counts are small whole numbers. Other top-k problems do not have this limit.

Which approach should you use?

The heap shows the pattern. The bucket version is the O(n) answer if the interviewer asks for something faster. (LeetCode 347.)

Q4K closest points to originMedium

You get points [x, y] on a flat plane. Return the k points closest to the origin (0, 0), measured in a straight line. The answer can be in any order.

Example 1
Input: points = [[1, 3], [-2, 2]], k = 1
Output: [[-2, 2]]
Explanation: The squared distances are 10 and 8, so (−2, 2) is closer.
Example 2
Input: points = [[3, 3], [5, -1], [-2, 4]], k = 2
Output: [[3, 3], [-2, 4]]
Explanation: The squared distances are 18, 26 and 20, so the two smallest are 18 and 20.
Stuck? Open one hint at a time
Hint 1
Do you need the real distance, or is the squared distance enough to compare? (You do not need a square root.)
Hint 2
For the k smallest, keep a max-heap of size k. The farthest of your k choices is on top, ready to be thrown out.
Show all 2 solutions with explanations
Approach 1Sort by distance
The idea

Sort all points by squared distance. Take the first k.

JavaScript
function kClosest(points, k) {  const dist = ([x, y]) => x * x + y * y;  return [...points].sort((a, b) => dist(a) - dist(b)).slice(0, k);}console.log(kClosest([[1, 3], [-2, 2]], 1));            // [ [ -2, 2 ] ]console.log(kClosest([[3, 3], [5, -1], [-2, 4]], 2));   // [ [ 3, 3 ], [ -2, 4 ] ]
How it works

It takes O(n log n) time. Squared distances stay whole numbers and give the same order as real distances.

Approach 2Max-heap of size k
The idea

Push each point into a max-heap ordered by distance. If it has more than k points, pop the farthest. The points left are the k closest.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function kClosest(points, k) {  const dist = ([x, y]) => x * x + y * y;  const heap = new Heap((a, b) => dist(b) - dist(a));   // max-heap by distance  for (const p of points) {    heap.push(p);    if (heap.size > k) heap.pop();  }  const result = [];  while (heap.size > 0) result.push(heap.pop());  return result.reverse();                              // closest first}console.log(kClosest([[1, 3], [-2, 2]], 1));            // [ [ -2, 2 ] ]console.log(kClosest([[3, 3], [5, -1], [-2, 4]], 2));   // [ [ 3, 3 ], [ -2, 4 ] ]
How it works

It takes O(n log k) time and O(k) space. It is the same pattern as the k-th largest, turned around. For the k smallest, use a max-heap of size k.

Which approach should you use?

Use the max-heap when k is small compared with n. Use the sort when you want the shortest code. Quickselect (see question 1) takes O(n) on average. (LeetCode 973.)

Q5Merge k sorted listsHard

You get an array of k sorted linked lists. Merge them into one sorted linked list and return its first node (the head).

Example 1
Input: lists = [[1, 4, 5], [1, 3, 4], [2, 6]]
Output: [1, 1, 2, 3, 4, 4, 5, 6]
Explanation: All eight values, in sorted order.
Example 2
Input: lists = []
Output: []
Explanation: There is nothing to merge.
Stuck? Open one hint at a time
Hint 1
Merging two sorted lists was lesson 37. How could you reuse it for k lists?
Hint 2
At any moment, the next answer is the smallest of k front nodes. What tool finds the smallest quickly?
Show all 3 solutions with explanations
Approach 1Collect everything and sort
The idea

Walk through every list and put all values in one array. Sort it, then build a new list.

JavaScript
class ListNode {  constructor(val, next = null) { this.val = val; this.next = next; }}function fromArray(values) {  const dummy = new ListNode(0);  let tail = dummy;  for (const v of values) { tail.next = new ListNode(v); tail = tail.next; }  return dummy.next;}function toArray(head) {  const out = [];  for (let cur = head; cur !== null; cur = cur.next) out.push(cur.val);  return out;}function mergeKLists(lists) {  const values = [];  for (let head of lists) {    for (; head !== null; head = head.next) values.push(head.val);  }  values.sort((a, b) => a - b);  return fromArray(values);}console.log(toArray(mergeKLists([fromArray([1, 4, 5]), fromArray([1, 3, 4]), fromArray([2, 6])]))); // [1, 1, 2, 3, 4, 4, 5, 6]console.log(toArray(mergeKLists([])));                                                               // []
How it works

It takes O(N log N) time, where N is the total number of nodes. It ignores that the lists are already sorted. It is a fine first answer.

Approach 2Min-heap of front nodes
The idea
  1. Push the first node of each list that is not empty into a min-heap ordered by value.
  2. Again and again, pop the smallest node, add it to the answer, and push the next node of the same list.
JavaScript
class ListNode {  constructor(val, next = null) { this.val = val; this.next = next; }}function fromArray(values) {  const dummy = new ListNode(0);  let tail = dummy;  for (const v of values) { tail.next = new ListNode(v); tail = tail.next; }  return dummy.next;}function toArray(head) {  const out = [];  for (let cur = head; cur !== null; cur = cur.next) out.push(cur.val);  return out;}class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function mergeKLists(lists) {  const heap = new Heap((a, b) => a.val - b.val);  for (const head of lists) if (head !== null) heap.push(head);  const dummy = new ListNode(0);  let tail = dummy;  while (heap.size > 0) {    const node = heap.pop();    tail.next = node;    tail = node;    if (node.next !== null) heap.push(node.next);  }  return dummy.next;}console.log(toArray(mergeKLists([fromArray([1, 4, 5]), fromArray([1, 3, 4]), fromArray([2, 6])]))); // [1, 1, 2, 3, 4, 4, 5, 6]console.log(toArray(mergeKLists([])));                                                               // []
How it works

It takes O(N log k) time and O(k) space. The heap size depends on the number of lists, not the number of nodes.

Approach 3Merge in pairs (divide and conquer)
The idea

Merge the lists two at a time. Then merge the results two at a time, and so on, until one list is left. It works like the rounds of a sports tournament.

JavaScript
class ListNode {  constructor(val, next = null) { this.val = val; this.next = next; }}function fromArray(values) {  const dummy = new ListNode(0);  let tail = dummy;  for (const v of values) { tail.next = new ListNode(v); tail = tail.next; }  return dummy.next;}function toArray(head) {  const out = [];  for (let cur = head; cur !== null; cur = cur.next) out.push(cur.val);  return out;}function mergeTwo(a, b) {  const dummy = new ListNode(0);  let tail = dummy;  while (a !== null && b !== null) {    if (a.val <= b.val) { tail.next = a; a = a.next; }    else { tail.next = b; b = b.next; }    tail = tail.next;  }  tail.next = a !== null ? a : b;  return dummy.next;}function mergeKLists(lists) {  if (lists.length === 0) return null;  let current = lists;  while (current.length > 1) {    const next = [];    for (let i = 0; i < current.length; i += 2) {      next.push(i + 1 < current.length ? mergeTwo(current[i], current[i + 1]) : current[i]);    }    current = next;  }  return current[0];}console.log(toArray(mergeKLists([fromArray([1, 4, 5]), fromArray([1, 3, 4]), fromArray([2, 6])]))); // [1, 1, 2, 3, 4, 4, 5, 6]console.log(toArray(mergeKLists([])));                                                               // []
How it works

There are about log₂ k rounds. Every round touches each node at most once. So the time is O(N log k). It relinks the existing nodes, and the only extra space is the small array of list heads, which is O(k). If you merged them one after another into a growing result, it would cost O(N·k) instead.

Which approach should you use?

The heap is the most common answer. Merging in pairs reaches the same time without a heap. (LeetCode 23.)

Q6Find median from data streamHard

Design a class with addNum(x) and findMedian(). The median is the middle value of the sorted numbers so far. If the count is even, it is the average of the two middle values.

Example 1
Input: addNum(5), findMedian(), addNum(2), findMedian(), addNum(8), findMedian()
Output: 5, 3.5, 5
Explanation: [5] has median 5. [2, 5] has the average of 2 and 5, which is 3.5. [2, 5, 8] has the middle value 5.
Stuck? Open one hint at a time
Hint 1
The median splits the data into a lower half and an upper half. Which item of each half do you need?
Hint 2
You need the biggest of the lower half and the smallest of the upper half. That means a max-heap and a min-heap.
Hint 3
After every add, keep the two sizes equal, or make the lower half one bigger.
Show all 2 solutions with explanations
Approach 1Keep a sorted array
The idea

Put each number in its sorted place. Read the median from the middle.

JavaScript
class MedianFinder {  constructor() { this.sorted = []; }  addNum(x) {    let i = 0;    while (i < this.sorted.length && this.sorted[i] < x) i++;    this.sorted.splice(i, 0, x);  }  findMedian() {    const n = this.sorted.length, mid = n >> 1;    return n % 2 === 1 ? this.sorted[mid] : (this.sorted[mid - 1] + this.sorted[mid]) / 2;  }}const m = new MedianFinder();const medians = [];for (const x of [5, 2, 8, 1, 9]) { m.addNum(x); medians.push(m.findMedian()); }console.log(medians); // [5, 3.5, 5, 3.5, 5]
How it works

Each add takes O(n) time (finding the place and moving items). Each median question takes O(1).

Approach 2Two heaps
The idea
  1. low is a max-heap with the smaller half. high is a min-heap with the larger half.
  2. To add a number: push it into low, then move the top of low to high. If high is now bigger, move its top back.
  3. To get the median: use the top of low if the count is odd. If the count is even, use the average of both tops.
JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}class MedianFinder {  constructor() {    this.low = new Heap((a, b) => b - a);    this.high = new Heap();  }  addNum(x) {    this.low.push(x);    this.high.push(this.low.pop());    if (this.high.size > this.low.size) this.low.push(this.high.pop());  }  findMedian() {    if (this.low.size > this.high.size) return this.low.peek();    return (this.low.peek() + this.high.peek()) / 2;  }}const m = new MedianFinder();const medians = [];for (const x of [5, 2, 8, 1, 9]) { m.addNum(x); medians.push(m.findMedian()); }console.log(medians); // [5, 3.5, 5, 3.5, 5]
How it works

Each add takes O(log n) time and each median takes O(1). Sending the new number through low into high makes sure that nothing in low is ever bigger than something in high. This works whichever half the number really belongs to.

Which approach should you use?

Use two heaps. The sorted-array version is a good way to start and explain the goal. (LeetCode 295.)

Q7Last stone weightEasy

Each turn, take the two heaviest stones and smash them together. If they weigh the same, both are destroyed. Otherwise, the lighter one is destroyed and the heavier one loses that much weight. Return the weight of the last stone, or 0 if no stone is left.

Example 1
Input: stones = [2, 7, 4, 1, 8, 1]
Output: 1
Explanation: Smash 8 and 7 to leave 1. Smash 4 and 2 to leave 2. Smash 2 and 1 to leave 1. Smash 1 and 1 to leave nothing. One stone of weight 1 remains.
Example 2
Input: stones = [1]
Output: 1
Explanation: There is only one stone, so there is nothing to smash.
Stuck? Open one hint at a time
Hint 1
Each round needs the two heaviest stones. What tool gives you the biggest item quickly?
Hint 2
After smashing, put the leftover back in the pile if it is not zero.
Show all 2 solutions with explanations
Approach 1Sort every round
The idea

While two or more stones remain, sort them and take out the two largest. Add back the difference if it is not zero.

JavaScript
function lastStoneWeight(stones) {  const pile = [...stones];  while (pile.length > 1) {    pile.sort((a, b) => a - b);    const y = pile.pop(), x = pile.pop();      // y is the heaviest, x is the second heaviest    if (y !== x) pile.push(y - x);  }  return pile.length === 1 ? pile[0] : 0;}console.log(lastStoneWeight([2, 7, 4, 1, 8, 1])); // 1console.log(lastStoneWeight([1]));                // 1console.log(lastStoneWeight([3, 3]));             // 0
How it works

It takes O(n² log n) time. There can be up to n rounds, and each round sorts up to n stones.

Approach 2Max-heap
The idea

Keep the stones in a max-heap. Pop two, then push back their difference if it is not zero.

JavaScript
class Heap {  constructor(compare = (a, b) => a - b) { this.data = []; this.compare = compare; }  get size() { return this.data.length; }  peek() { return this.data[0]; }  push(x) {    const d = this.data; d.push(x);    for (let i = d.length - 1; i > 0; ) {      const p = (i - 1) >> 1;      if (this.compare(d[i], d[p]) >= 0) break;      [d[i], d[p]] = [d[p], d[i]]; i = p;    }  }  pop() {    const d = this.data;    if (d.length === 0) return undefined;    const top = d[0], last = d.pop();    if (d.length) {      d[0] = last;      for (let i = 0; ; ) {        const l = 2 * i + 1, r = l + 1;        let b = i;        if (l < d.length && this.compare(d[l], d[b]) < 0) b = l;        if (r < d.length && this.compare(d[r], d[b]) < 0) b = r;        if (b === i) break;        [d[i], d[b]] = [d[b], d[i]]; i = b;      }    }    return top;  }}function lastStoneWeight(stones) {  const heap = new Heap((a, b) => b - a);       // max-heap  for (const s of stones) heap.push(s);  while (heap.size > 1) {    const y = heap.pop(), x = heap.pop();    if (y !== x) heap.push(y - x);  }  return heap.size === 1 ? heap.peek() : 0;}console.log(lastStoneWeight([2, 7, 4, 1, 8, 1])); // 1console.log(lastStoneWeight([1]));                // 1console.log(lastStoneWeight([3, 3]));             // 0
How it works

It takes O(n log n) time. Each round does a fixed number of heap steps and removes at least one stone.

Which approach should you use?

Use the max-heap. This is the classic "keep taking the largest" problem. (LeetCode 1046.)

Make it stick

Make it stick — close this page, then
  1. Say the heap rule, and give the index formulas for the parent and the children of an item.
  2. Explain push (sift up) and pop (sift down), and why each is O(log n).
  3. Say why the k-th largest uses a min-heap of size k, and what it costs.
  4. Explain how two heaps give the median of a stream in O(log n) per number.

Revisit schedule: solve two of the questions again tomorrow, again in 3 days, and once more in a week — without looking at the answers. Each time it will feel easier. That is the memory forming.

What's next

A heap picks the best item of what is left, again and again. Lesson 47 uses the same idea in a different place: greedy algorithms. There you take the choice that looks best right now, and you prove it never hurts later.

Practise on LeetCode

Real interview problems that use exactly what this lesson teaches. Solve them in JavaScript after the practice questions above.

Published by Jainik · Last updated · How these lessons are made

Keep learning

Part of the study guide

Related series