Top-k with a size-k heap

41 s · title card, the worked example step by step, then a recap. The captions carry the explanation; there is no narration. Download the MP4.

01The intuition

Start with what you would do by hand

Numbers arrive one at a time: 4, 5, 8, 2, 3, 5, 10, 9. After each arrival you want to know the third largest so far. By hand you would keep a sorted list of everything, insert each newcomer in its place, and read the third entry from the end. Inserting into a sorted list of m numbers means shifting about m/2 of them, so after a million numbers each insert costs about half a million moves. And you keep all million numbers although only three matter.

The first idea: only keep the three biggest seen so far. A newcomer either beats the smallest of the three, and replaces it, or does not, and is thrown away. Now the list has three entries and inserting is cheap. The question is how to do this when k is 1,000 instead of 3.

Build the structure, one piece at a time

Piece one: a min-heap. A heap is a collection where you can always see its smallest element at once, add a new element, and remove the smallest, each of the last two costing about log₂(size) steps. Internally it is a tree where every node is smaller than its children; the root is therefore the minimum. For now, treat it as a bag with a cheap "smallest" button.

Piece two: cap the bag at k elements. Push the first k numbers in. From then on, compare each newcomer with the smallest in the bag. If the newcomer is bigger, remove the smallest and push the newcomer. Otherwise ignore it.

Piece three: read the answer. The bag holds the k largest numbers seen so far, and its smallest element is therefore the k-th largest overall.

Why a min-heap for the largest elements: the one element we must compare against and remove is the smallest of the kept ones, so that is the element the structure must expose.

Why it is correct

The invariant: after any number of arrivals, the bag contains exactly the k largest numbers seen so far. True after the first k. Consider a newcomer x and the bag's smallest s. If x ≤ s, then at least k numbers are at least as large as x, so x is not among the k largest, and the bag is unchanged correctly. If x > s, then s is beaten by x and by the other k − 1 bag members, so s drops out and x joins, which is exactly the swap we perform.

min(bag) = k-th largest of everything seen, because exactly k − 1 seen numbers exceed it

Why it is fast

Each arrival does one comparison and at most one remove and one push on a heap of size k, each O(log k). For n arrivals: O(n log k), and only k numbers are stored. With k = 3 that is about 2 steps per number instead of half a million moves in the sorted list. Sorting everything would be O(n log n) and needs all n in memory, which fails for a stream.

How to recognise it in an interview

"The k largest", "k closest", "k most frequent", "kth largest in a stream", especially when k is much smaller than n or the data arrives continuously. Say: "I will keep a min-heap of size k; the root is the kth largest, and a newcomer replaces the root only if it is bigger." Kth Largest Element in a Stream is this object with an add() method. K Closest Points to Origin uses a max-heap of size k on distance, the mirror image. Top K Frequent Elements first counts, then runs the same heap on the counts.

02The idea in one line

Keep a min-heap of the k best; a newcomer replaces the root only if it beats it.

Reach for it when: Kth largest, k closest points, top-k frequent, kth largest in a stream.

Cost: O(n log k) time.

03The walk-through, step by step

The video works one example: Kth largest (k=3) stream on [4,5,8,2,3,5,10,9]. These are its captions in order, so you can follow without the video or check a step you missed.

1. Keep a min-heap of size k: its root is the kth largest
2. Heap has room: push 4
3. Heap has room: push 5
4. Heap has room: push 8
5. 2 ≤ root 4: it cannot be in the top 3, skip
6. 3 ≤ root 4: it cannot be in the top 3, skip
7. 5 beats the root 4: replace the root and sift down
8. 10 beats the root 5: replace the root and sift down
9. 9 beats the root 5: replace the root and sift down
10. After the stream the root is the 3rd largest: 8
11. Each item costs one heap op of size k: O(n log k)

04Asked at Google

7 Google-tagged LeetCode problems use this pattern. Frequency is LeetCode's Google-tag score (higher means reported more often); "last 6 months" marks problems reported in the most recent window of the export.

05Related patterns

Other patterns in Heap: Two heaps (running median), K-way merge, Greedy scheduling with a heap, Lazy deletion heap, Pop-k-then-rebuild (frequency greedy).