Top-K / Heap
Why top-k heaps exist — the story
Imagine you are watching a live scoreboard with millions of scores, but the product manager only cares about the top 10. Sorting every score every time would be silly: rank 11 through rank 1,000,000 do not matter. The Top-K heap pattern exists because you can keep a tiny "VIP room" of size k instead of arranging the whole crowd. Every new candidate enters the room briefly; if the room is too full, you kick out the worst person currently inside. At the end, the room contains exactly the best k candidates.
Use a concrete example: find the 3 largest numbers in [7, 1, 9, 3, 10, 2, 8]. Keep a min-heap of size 3. Read 7,1,9 → heap has {1,7,9}. Read 3 → heap becomes {1,3,7,9}, too large, so remove the smallest 1; room is {3,7,9}. Read 10 → remove 3; room is {7,9,10}. Read 2 → remove 2 immediately. Read 8 → remove 7; final room is {8,9,10}. Notice the heap root is not the answer to "largest"; it is the weakest of the winners, the one easiest to evict.
The brute-force walk-through is full sorting: arrange everyone, then take the first k. Can we do better? Yes, because the exact order of the rejected n-k items does not matter. Keep only the current boundary set.
That opposite-polarity idea is the whole trick. For k largest, use a min-heap. For k smallest, use a max-heap. For k most frequent, use a min-heap ordered by frequency. The heap is not there to sort everything; it is there to maintain a moving boundary between "still in the top k" and "not good enough."
For k largest, use a min-heap (its root is the worst-of-the-best, cheap to evict). For k smallest, use a max-heap. Polarity is always the opposite of the goal.
When to use it — boundary items without sorting
Top-K heap is the interview pattern for "I need a small elite set from a large input." It is especially strong when k is much smaller than n, when the input is streaming, or when the final order of the selected items does not matter.
Recognize by
- "k largest", "k smallest", "top k", "bottom k"
- "k most frequent", "k least frequent", "most common words"
- "k closest to origin" or "nearest k elements" where the ranking key is computed from each item
- "kth largest in a stream" where values arrive over time and you must answer after each add
- "return any order" or "relative order among the top k does not matter"
- constraints like
n = 10^5andksmall enough that O(n log k) is attractive
When NOT to use it
You need the k-th value once and don't care about the other k−1 boundary items. Quickselect is O(n) average and beats the heap's O(n log k). Use a heap when the input arrives as a stream (you can't Quickselect it) or when you need all k boundary items.
Also avoid a heap when you need the complete sorted order of all values; then sorting is simpler and honest. If the key range is tiny or bounded, bucket sort can beat the heap: for frequencies from 0 to n, buckets often give O(n). If k is close to n, O(n log k) becomes close to O(n log n), so the benefit shrinks.
How to use it — template
PriorityQueue<Item> heap = new PriorityQueue<>((a, b) -> compareWorstFirst(a, b));
for (Item item : items) {
heap.offer(item);
if (heap.size() > k) {
heap.poll(); // remove the worst of the kept items
}
}
List<Item> answer = new ArrayList<>();
while (!heap.isEmpty()) {
answer.add(heap.poll());
}
return answer;Fill in compareWorstFirst based on what "worst among the winners" means. For k largest, the smallest value is worst, so Java's default min-heap works. For k smallest, the largest value is worst, so reverse the comparator. For k most frequent, the lowest frequency is worst, so order by freq.get(x). The loop invariant is simple: after processing any prefix of the input, the heap contains the best k items from that prefix, or all items if fewer than k have appeared.
Never grows beyond k, never shrinks. The "grow then trim" one-liner offer; if (size > k) poll; maintains this exactly. If size drifts, the algorithm silently returns the wrong top-k.
Invariant — the heap root is the weakest of the winners. For k-largest with a min-heap, the root is the smallest current winner — it's the one to compare a new candidate against. If candidate > root, evict root and add candidate. If candidate <= root, the candidate can never be in the top-k, so discard. This one comparison per element is what makes it O(n log k) instead of O(n log n).
When NOT to use it
- You need the top k in sorted order relative to each other — the heap doesn't guarantee internal order, only that the root is the worst-of-the-winners. Extract with successive polls if order matters.
- k is very close to n (say k = n/2) — the heap approach becomes O(n log(n/2)) ≈ O(n log n), same as full sort. Just sort.
- You need to answer many top-k queries with different k on the same data — use quickselect once to partition, or sort once and slice.
- Elements arrive as a stream and you need to remove arbitrary items later — a plain heap can't remove non-root elements in O(log n). Use
TreeSetor an indexed heap. - k = 1 — you don't need a heap. Just track a running max / min with a scalar. Simpler, faster, no allocation.
- The comparator involves floating-point tolerance — heap operations don't tolerate
NaNor unstable comparators. Wrap in integer keys where possible.
Traps & gotchas — the 5 that fail candidates on interview day
For k largest, use a min-heap (worst-of-winners at root). Candidates instinctively reach for max-heap and end up popping the wrong elements. Say the polarity rule aloud before writing: "k largest → min-heap; k smallest → max-heap."
Trap 2 — Adding then checking, vs checking then adding. The clean template is offer(item); if (heap.size() > k) poll();. Some candidates write if (heap.size() < k) offer(item); else if (item beats root) { poll(); offer(item); }. Both work but the second is bug-prone (need to compare item to root correctly). Prefer the first.
For primitives it's fine: new PriorityQueue<>(Collections.reverseOrder()) gives a max-heap of Integers. For custom classes, wrap the natural comparator: new PriorityQueue<>((a, b) -> b.value - a.value) — and beware int subtraction overflow when values span Integer.MIN_VALUE. Use Integer.compare(b.value, a.value).
After the loop, poll() returns the worst of the winners first. If the interviewer wants "top k in descending order", you'll need to reverse the extracted list — or use a max-heap from the start with size n and poll k times (O(n + k log n), often faster).
Java's PriorityQueue<Integer> boxes every int into an Integer. For n = 10⁷, that's 10⁷ heap allocations for the boxed keys alone. In performance-critical contexts, use a primitive heap (from Eclipse Collections or write your own on a long[]). Interviews rarely test this, but staff-level candidates should mention it.
History — Williams's heap, 1964
The binary heap was invented by J.W.J. Williams in 1964 as the core data structure of heapsort. Williams's paper introduced both the heap-invariant maintenance operations (sift-up, sift-down) and the array-based tree representation using index arithmetic — parent at (i-1)/2, children at 2i+1 and 2i+2. That representation lets a heap live entirely in a contiguous int[] with zero pointer overhead, which is why every language's standard library ships a heap-backed PriorityQueue.
The k-way merge algorithm — used in every external sort since von Neumann's 1945 tape-sort memo — is built on the same heap idea, with the heap holding one live iterator per input stream. Modern implementations include Lucene's segment merging, DuckDB's external sort, and Kafka's log compaction.
The specific Top-K-on-a-stream pattern was formalized in the NIST Text REtrieval Conference (TREC) benchmarks in the late 1990s as "top-K retrieval," and later applied to trending / recommendation systems at Google, Facebook, and Twitter. Every "top trending topic" list you see is a Top-K heap in disguise.
Canonical problem walkthrough — K Closest Points to Origin
Problem (↗ LeetCode): Given an array of points on a 2D plane, return the k closest points to the origin. Distance is Euclidean; the order of the returned points doesn't matter.
Approach 1 — Full sort by distance
int[][] kClosestSort(int[][] points, int k) {
Arrays.sort(points, (a, b) -> (a[0]*a[0] + a[1]*a[1]) - (b[0]*b[0] + b[1]*b[1]));
return Arrays.copyOf(points, k);
}Complexity: O(n log n) time, O(1) extra space (in-place sort). Correct, one line, but does more work than needed when k << n.
Approach 2 — Max-heap of size k
Keep a heap of size k where the root is the farthest point currently in the top-k. Each new point that's closer than the root evicts it.
int[][] kClosest(int[][] points, int k) {
// Max-heap by squared distance (avoid sqrt — it's monotonic so it doesn't change order).
PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> {
int da = a[0]*a[0] + a[1]*a[1], db = b[0]*b[0] + b[1]*b[1];
return db - da; // farthest at root
});
for (int[] p : points) {
heap.offer(p);
if (heap.size() > k) heap.poll(); // evict the farthest
}
int[][] res = new int[k][2];
for (int i = 0; i < k; i++) res[i] = heap.poll();
return res;
}Complexity: O(n log k) time, O(k) space.
Why this is faster when k is small: each of the n iterations does O(log k) work instead of the O(log n) of full sort. For n = 10⁵, k = 10: sort is ~1.7·10⁶ ops, heap is ~3.3·10⁵ ops — 5× faster. The gap widens as k shrinks.
Approach 3 — Quickselect (Hoare partition)
Partition the array around a pivot's distance; recurse on the side that contains rank k. Average O(n).
int[][] kClosestQuickselect(int[][] points, int k) {
quickselect(points, 0, points.length - 1, k);
return Arrays.copyOf(points, k);
}
void quickselect(int[][] a, int lo, int hi, int k) {
if (lo >= hi) return;
int pivotIdx = lo + (int)(Math.random() * (hi - lo + 1));
int pivotDist = dist(a[pivotIdx]);
swap(a, pivotIdx, hi);
int store = lo;
for (int i = lo; i < hi; i++)
if (dist(a[i]) < pivotDist) swap(a, i, store++);
swap(a, store, hi);
if (store == k) return;
else if (store < k) quickselect(a, store + 1, hi, k);
else quickselect(a, lo, store - 1, k);
}
int dist(int[] p) { return p[0]*p[0] + p[1]*p[1]; }
void swap(int[][] a, int i, int j) { int[] t = a[i]; a[i] = a[j]; a[j] = t; }Complexity: O(n) average, O(n²) worst case (defended by random pivot).
When to reach for it: contest programming or when you know k is close to n/2. Interview trade-off — quickselect is faster asymptotically but more code, and interviewers occasionally distrust the "expected O(n)" claim.
Complexity ladder
| Approach | Time | Space | When |
|---|---|---|---|
| Sort by distance | O(n log n) | O(1) | Simplest; equally good when k ~ n |
| Max-heap of size k | O(n log k) | O(k) | Interview default when k ≪ n |
| Quickselect | O(n) avg / O(n²) worst | O(1) in-place | Contest / follow-up "make it O(n)" |
Kth Largest / Top K Frequent Medium
↗ LeetCode: Top K Frequent Elements
Problem
Return the k most frequent elements of an array (or, in the sibling, the kth largest). The order among the top-k doesn't matter.
Constraints: 1 ≤ n ≤ 10⁵; 1 ≤ k ≤ #distinct; beat a full O(n log n) sort.
Example 1: nums = [1,1,1,2,2,3], k = 2 → [1,2].
Example 2: nums = [1], k = 1 → [1].
Solution — brute force
First count frequencies. The brute-force way is then to sort all distinct values by decreasing frequency and take the first k.
freq = count each number
items = all distinct numbers
sort items by freq descending
return first k itemsint[] topKFrequentBrute(int[] nums, int k) {
Map<Integer,Integer> freq = new HashMap<>();
for (int x : nums) freq.merge(x, 1, Integer::sum);
List<Integer> keys = new ArrayList<>(freq.keySet());
keys.sort((a, b) -> freq.get(b) - freq.get(a));
int[] res = new int[k];
for (int i = 0; i < k; i++) res[i] = keys.get(i);
return res;
}This is perfectly acceptable as a baseline: O(n) to count, O(m log m) to sort m distinct values, and O(m) space. It becomes wasteful when m is large and k is small, because sorting rank 5000 versus rank 5001 does not help you return the top 10. The heap version avoids fully ordering the losers.
O(n + m log m) time, O(m) space, where m is the number of distinct values — wasteful when k is tiny.
Solution — optimized
Size-k min-heap; evict the smallest whenever the heap exceeds k.
The heap stores distinct numbers, not array indices and not (number, frequency) pairs. The comparator looks up each number's count in the frequency map. Because this is a min-heap by frequency, poll() removes the least frequent number among the currently kept candidates. After all distinct values have been processed, only the k most frequent can survive.
After each key from freq.keySet() is processed, pq contains the k highest-frequency keys seen so far. If more than k keys have been offered, the lowest-frequency survivor has already been evicted.
The optimized Java counts once, then keeps a heap of only k surviving keys. The root is deliberately the weakest survivor, so overflow beyond size k evicts the least useful candidate immediately.
int[] topKFrequent(int[] nums, int k) {
Map<Integer,Integer> freq = new HashMap<>();
for (int x : nums) freq.merge(x, 1, Integer::sum);
PriorityQueue<Integer> pq = // min-heap by frequency
new PriorityQueue<>((a, b) -> freq.get(a) - freq.get(b));
for (int key : freq.keySet()) {
pq.offer(key);
if (pq.size() > k) pq.poll(); // drop least frequent
}
int[] res = new int[k];
for (int i = k - 1; i >= 0; i--) res[i] = pq.poll();
return res;
}nums=[1,1,1,2,2,3], k=2.
Space Complexity
O(n), because the frequency map may store all n values when every value is distinct; the heap itself stores only O(k).
More precisely, counting costs O(n), and heap operations cost O(m log k), where m is the number of distinct values. The frequency map can hold up to O(n) entries. The heap itself holds O(k) entries, but the total auxiliary space is still O(n) because of the map.
Wrong heap polarity. Example: the k-th largest with a max-heap of all n elements costs O(n log n) — wasteful. A min-heap of size k evicts the smallest; the root is your answer in O(n log k).
First, I'd count each number because frequency is the ranking key. The brute force is to sort all distinct numbers by frequency and take the first k, which is O(n + m log m). Since we only need k winners, I'll keep a min-heap of size k ordered by frequency and evict the least frequent whenever the heap grows too large. That gives O(n + m log k) time and O(n) space, and the output order can be arbitrary because the problem allows it.
K Closest Points to Origin is identical with a max-heap of size k keyed on distance; bucket sort by frequency gives O(n) when the key range is bounded.
Learning notes
- Why count first? — frequency is the ranking key, so the heap cannot compare numbers until counts exist.
- Why store just the key in the heap? — the
freqmap already holds the count; duplicating pairs is unnecessary here. - Why a min-heap for "most frequent"? — the least frequent survivor is the next one to evict.
- Why
if (pq.size() > k) pq.poll()? — allowing one temporary extra item lets each candidate compete before eviction. - Why fill
resfrom right to left? — repeatedpoll()returns lower-frequency survivors first. - Why not sort all keys? — sorting orders every loser too; the heap only maintains the top-k boundary.
- Why mention
Integer.compare? — subtraction comparators can overflow when values or counts are extreme.
Choosing between heap, sort, bucket, and quickselect
A useful interview move is to name the trade-off instead of pretending the heap is always best. Full sorting is shortest to code and good when n is small or when the final order matters. A size-k heap is best when k is small, the stream may not fit in memory, or you need to update answers incrementally. Bucket sort is excellent for frequencies because frequency is an integer from 1 to n; create buckets where bucket f stores all values with frequency f, then read from high frequency down until k values are collected. Quickselect is attractive for one-shot selection when you can mutate the array of candidates and do not need streaming behavior.
For Top K Frequent specifically, the heap solution is a reliable general answer because it does not depend on a small value range and is easy to adapt to words, points, pairs, or custom objects. If the interviewer asks for strict O(n), discuss buckets after presenting the heap.
Comparator details in Java
Java's PriorityQueue is a min-heap according to the comparator. That means the "smallest" item by comparator comes out first. For top-k largest, you intentionally define "smallest" as the item you are most willing to discard. In the existing code, (a, b) -> freq.get(a) - freq.get(b) puts lower-frequency numbers first.
In production code, Integer.compare(freq.get(a), freq.get(b)) avoids overflow. The chapter keeps the compile-tested solution unchanged, but in an interview you can mention Integer.compare if counts or values can be extreme. For tie-breaking, add a secondary comparison only if the problem requires deterministic ordering; otherwise, extra tie logic is unnecessary.
Heap polarity exercises
When you are under interview pressure, decide heap polarity with one sentence: the root should be the next thing I would throw away. If you want k largest numbers, the smallest kept number is the next thing to throw away, so the root must be small: min-heap. If you want k smallest numbers, the largest kept number is the next thing to throw away, so the root must be large: max-heap. If you want k closest points, the farthest kept point is the next thing to throw away, so use a max-heap by distance. If you want k most frequent words, the least frequent kept word is the next thing to throw away, so use a min-heap by frequency.
Try this quick table before coding:
| Goal | Worst survivor | Heap at root |
|---|---|---|
| k largest values | smallest value among the kept k | min-heap |
| k smallest values | largest value among the kept k | max-heap |
| k closest points | largest distance among the kept k | max-heap |
| k most frequent | smallest frequency among the kept k | min-heap |
| k least frequent | largest frequency among the kept k | max-heap |
This table prevents the most common reversal bug. The heap is not a trophy case where the best item sits on top. It is a door guard; the root is the easiest item to remove when a better candidate arrives.
Streaming vs batch thinking
Top-K heap is the natural streaming answer because it never needs to see the future. After processing the first 100 numbers, the heap already contains the top k among those 100. After the 101st number, one offer and maybe one poll restores the same invariant. That is why Kth Largest Element in a Stream stores the heap as object state and returns pq.peek() after each add.
Quickselect cannot do that. Quickselect needs a closed array to partition; when a new number arrives, the previous partition work may no longer line up with the new rank. If an interviewer asks, "What if numbers arrive continuously?" switch from Quickselect to heap immediately. If they ask, "What if this is one offline query and k is large?" Quickselect becomes attractive.
Tie handling and output order
Many top-k problems say "return the answer in any order." That permission matters. The heap output order is usually from worst survivor to best survivor because repeated poll() removes the comparator-minimum. The existing solution fills the result from right to left so higher-frequency values tend to appear earlier, but the LeetCode problem does not require a sorted top-k list.
If a problem requires deterministic ties, bake the tie-breaker into the comparator. For words, you might sort by frequency ascending but reverse lexicographic order for the min-heap so that the lexicographically worse word is evicted first. For points, equal distances might not need tie-breaking at all. Never add tie rules that the prompt does not ask for; they make comparators harder to reason about and can accidentally evict the wrong item.
Testing checklist
Run these small cases mentally before trusting the code:
| Case | Expected lesson |
|---|---|
nums=[1], k=1 | heap and result handle the smallest input |
nums=[1,2,3], k=3 | no distinct key should be evicted permanently |
nums=[1,1,2,2], k=1 | ties are acceptable unless prompt specifies order |
nums=[1,1,1,2,3], k=2 | a low-frequency late key should be offered then immediately evicted |
| all values distinct | any k values are valid for top-k frequent because all frequencies tie |
Same pattern, new tweaks
"Keep only the best k in a size-k heap of the opposite polarity" transfers straight across:
| Variation | The one thing that changes | Time |
|---|---|---|
| K Closest Points to Origin | a max-heap of size k keyed on squared distance; the farthest current winner is easiest to evict | O(n log k) |
| Top K Frequent Elements | heap keyed on frequency, or bucket by frequency for O(n) when you want the linear variant | O(n log k) |
| Kth Smallest Element in a Sorted Matrix | treat each row as a sorted stream and use a min-heap for k-way merge | O(k log n) |
| Kth Largest Element in a Stream | keep a size-k min-heap alive across add calls; root is always the running kth largest | O(log k) per add |
| Reorganize String | use a max-heap because you repeatedly need the currently most frequent remaining character, not a bounded top-k set | O(n log alphabet) |