Heap of the best k

When to use it

Use it when only the k largest or k smallest of many values matter, and every other value can be thrown away as soon as k others are better.

It belongs to the trees group: Let the shape of the data do the work: the branches of a tree, the letters of a trie, the top of a heap.

Watch it work

In a fishing contest your score is the total weight of your k heaviest fish. These are the weights of the fish you caught, in the order you caught them. Return your score.

Input
weights = [4, 9, 2, 7, 3, 8]
k = 3

Output
24

The method

1. Start with an empty min-heap: its top is the smallest it holds
2. Push each weight onto the heap
3. More than k in the heap: pop the top, the smallest of them
4. The heap now holds the k heaviest: add them up

The run, one moment at a time

1. Step 1. The heap is empty. It will keep at most 3 weights, and its top will always be the lightest of them.

   weights: 4 9 2 7 3 8
   heap: empty
   

2. Step 2. Push 4. The heap holds 1 of 3 weights, so nothing has to leave. Its top, 4, is the lightest kept.

   weights: 4 9 2 7 3 8
   heap: 4  (top at 0)
   

3. Step 2. Push 9. The heap holds 2 of 3 weights, so nothing has to leave. Its top, 4, is the lightest kept.

   weights: 4 9 2 7 3 8
   heap: 4 9  (top at 0)
   

4. Step 2. Push 2. The heap now holds 3 weights and is full. Its top, 2, is the lightest of them.

   weights: 4 9 2 7 3 8
   heap: 2 9 4  (top at 0)
   

5. Step 2. Push 7. The heap holds 4 weights, one more than k. Its top, 2, is the lightest of the 4.

   weights: 4 9 2 7 3 8
   heap: 2 7 4 9  (top at 0)
   

6. Step 3. Pop the top, 2. The 3 fish still in the heap are all heavier, so 2 can never be among the 3 heaviest.

   weights: 4 9 2 7 3 8
   heap: 4 7 9  (top at 0)
   

7. Step 2. Push 3. The heap holds 4 weights, one more than k. Its top, 3, is the lightest of the 4.

   weights: 4 9 2 7 3 8
   heap: 3 4 9 7  (top at 0)
   

8. Step 3. Pop the top, 3: the fish pushed a moment ago leaves again. The 3 fish in the heap are all heavier, so it can never count.

   weights: 4 9 2 7 3 8
   heap: 4 7 9  (top at 0)
   

9. Step 2. Push 8. The heap holds 4 weights, one more than k. Its top, 4, is the lightest of the 4.

   weights: 4 9 2 7 3 8
   heap: 4 7 9 8  (top at 0)
   

10. Step 3. Pop the top, 4. The 3 fish still in the heap are all heavier, so 4 can never be among the 3 heaviest.

    weights: 4 9 2 7 3 8
    heap: 7 8 9  (top at 0)
    

11. Step 4. All 6 fish were looked at. The 3 weights left in the heap, 7, 8 and 9, are the 3 heaviest. Their total is 24.

    weights: 4 9 2 7 3 8
    heap: 7 8 9  (top at 0)
    total: 24
    

The code

The numbered comments are the steps of the method.

// JavaScript has no built-in heap, so these two helpers keep an array
// in min-heap order by hand: no value is smaller than its parent, so
// the smallest is at index 0. The parent of index i is at
// (i - 1) >>> 1, which is (i - 1) / 2 rounded down, and its two
// children are at 2 * i + 1 and 2 * i + 2.
function heapPush(heap, x) {
  // The new value starts at the end and swaps with its parent while the
  // parent is larger.
  heap.push(x);
  let i = heap.length - 1;
  while (i) {
    const p = (i - 1) >>> 1;
    if (heap[p] <= heap[i]) break;
    [heap[p], heap[i]] = [heap[i], heap[p]];
    i = p;
  }
}
function heapPop(heap) {
  // The smallest value leaves. The last value takes its place at
  // index 0, then swaps with its smaller child while that child is
  // smaller.
  const smallest = heap[0],
    last = heap.pop();
  if (!heap.length) return smallest;
  heap[0] = last;
  let i = 0;
  while (true) {
    let child = 2 * i + 1;
    if (child >= heap.length) break;
    if (child + 1 < heap.length && heap[child + 1] < heap[child]) {
      child++;
    }
    if (heap[i] <= heap[child]) break;
    [heap[i], heap[child]] = [heap[child], heap[i]];
    i = child;
  }
  return smallest;
}

function solve(weights, k) {
  // 1. Start with an empty min-heap: its top is the smallest it holds
  const heap = [];
  for (const weight of weights) {
    // 2. Push each weight onto the heap
    heapPush(heap, weight);
    // 3. More than k in the heap: pop the top, the smallest of them
    if (heap.length > k) heapPop(heap);
  }
  // 4. The heap now holds the k heaviest: add them up
  let total = 0;
  for (const weight of heap) total += weight;
  return total;
}

import heapq


def solve(weights, k):
    # 1. Start with an empty min-heap: its top is the smallest it holds
    # heapq keeps a plain list in heap order: heap[0] is its smallest.
    heap = []
    for weight in weights:
        # 2. Push each weight onto the heap
        heapq.heappush(heap, weight)
        # 3. More than k in the heap: pop the top, the smallest of them
        if len(heap) > k:
            heapq.heappop(heap)
    # 4. The heap now holds the k heaviest: add them up
    total = 0
    for weight in heap:
        total += weight
    return total

How to spot it

Asked: Return the k largest values, or the kth largest

Only k values can matter. A min-heap, a heap with its smallest on top, keeps them: the top is the weakest of the k.

Asked: The k closest, cheapest or smallest of many items

Keep k items in a max-heap, a heap with its largest on top. The top is the worst of the k kept, so a better item pushes it out.

Given: Values arrive one at a time, and the best k are wanted after each

Updating a heap of k takes about log k steps. Sorting everything again after each new value is far slower.

Limits: n is very large and k is small

Sorting costs about n log n steps and keeps all n values. A heap of k costs about n log k steps and keeps only k.

Looks like it, but is not

• You must take the most urgent item again and again, and no item may be thrown away. Use Heap as a priority queue. There the heap holds everything that waits and hands the items out one by one. Here it is capped at k and the rest is dropped.
• You must report the middle value of the numbers, not the k largest. Use Two heaps. The middle moves as numbers arrive, so no fixed k fits. Two heaps keep a smaller half and a larger half around it.
• k is the length of a stretch of neighbours in a list, and you want the largest value in each stretch. Use Monotonic queue. There a value must leave when it falls out of the stretch, whatever its size. A heap of the best k only ever drops its top.

Cost and mistakes

Each of the n values is pushed once and popped at most once. The heap never holds more than k + 1 values, so each of those takes about log k steps: O(n log k) time. Extra space is the k values kept, O(k).

Watch out

• To keep the k largest, use a min-heap, not a max-heap. The top must be the weakest of the values kept, because that is the one to throw away.
• A heap is not a sorted list. Only its top is in a known place, so heap[1] is not always the second smallest. Sort the k values at the end if you need them in order.

In the 150

3 problems of the 150 use it. Naming them gives their method away, so the list is folded.