Episode 8: "Heap or Sort? The Question Behind Every Top-K Problem"
Understand when heaps beat sorting and how min-heaps, max-heaps, and the size-K trick solve Top-K interview problems.
The Coding Interview Pattern Series covers the ten patterns that handle the large majority of coding interview problems. Episodes 1 through 7 covered hash maps, sliding windows, binary search and two pointers, graphs, dynamic programming, fast and slow pointers, and the pattern diagnosis drill. This is Episode 8: heaps, the pattern behind every top-K problem and any situation that repeatedly needs the current minimum or maximum from a changing collection.
Here is a question that appears in many forms across coding interviews. You have a million numbers and you want the largest hundred.
What do you do?
The first instinct is to sort everything and take the top hundred. That works. It is also doing approximately ten million operations (a million times log of a million) to produce a hundred results. You sorted nine hundred and ninety-nine thousand nine hundred numbers you will never use.
A heap does the same job in a fraction of the work. Keep a min-heap of exactly a hundred numbers. Walk through the million.
For each number, push it into the heap; if the heap now has a hundred and one elements, pop the smallest (which is now the smallest of the top hundred and one, not the overall answer you want). When you finish, the heap contains exactly the hundred largest numbers.
Total work: a million pushes and at most a million pops, each costing log of a hundred (about seven operations). Approximately fourteen million operations instead of ten million?
Actually less: each heap operation is O(log K) where K is a hundred, not log of a million. That is seven versus twenty. Across a million elements, the heap approach is nearly three times faster.
For small K that difference is dramatic.
For K equal to n, sorting wins. The crossover point is roughly when K is small relative to n, which covers the vast majority of “top K” problems in interviews.
This episode teaches the heap pattern completely: what a heap is and why it works, the size-K trick that is the core move for top-K problems, four problems that show the pattern in different disguises, and the two-heap technique that solves the median-of-a-stream problem.
What a Heap Is and What it Guarantees
A heap is a data structure that makes one operation very fast: finding the minimum (or maximum) element of a changing collection in O(1) time. Insertions and deletions are O(log n). You do not get efficient access to arbitrary elements, only to the extreme one.
In Python, heapq implements a min-heap: the smallest element is always at position zero. There is no built-in max-heap; the standard approach is to negate values, turning a min-heap into an effective max-heap.
import heapq
h = []
heapq.heappush(h, 5)
heapq.heappush(h, 3)
heapq.heappush(h, 8)
heapq.heappush(h, 1)
print(h[0]) # -> 1, the minimum, always at position 0
heapq.heappop(h) # -> 1, removes and returns the minimum
print(h[0]) # -> 3, the new minimum
The heap is built on a binary tree structure with one invariant: every parent is smaller than its children (min-heap). This invariant guarantees that the minimum is always at the root (position zero) and that inserting or removing takes O(log n) time to restore the invariant by bubbling up or down.
You do not need to understand the tree structure to use a heap effectively. You need to understand the guarantee: the minimum element is always instantly accessible, and adding or removing elements costs O(log n).
The Size-K trick: the Core Move for top-K Problems
The size-K trick is the single most important heap technique, and it is counterintuitive enough to be worth stating clearly before any examples.
To find the K largest elements, use a min-heap of size K.
The logic: the min-heap of size K always contains the K largest elements seen so far. The element at the top (the minimum of the heap) is the smallest of the current top K. When a new element arrives, compare it to the heap minimum. If it is larger, it belongs in the top K, so push it and pop the previous minimum. If it is smaller or equal, it does not belong in the top K, so ignore it (or push and immediately pop). After processing all elements, the heap contains exactly the K largest.
import heapq
def top_k_largest(nums, k):
heap = []
for n in nums:
heapq.heappush(heap, n)
if len(heap) > k:
heapq.heappop(heap) # evict the smallest; keep only the k largest
return sorted(heap, reverse=True)
top_k_largest([3,1,5,2,8,7,4], 3) # -> [8, 7, 5]
The part that confuses people: you use a min-heap to track the largest elements. The logic is: the element at the top of the min-heap is the smallest of the current top K candidates. It is the one most likely to be displaced by a larger newcomer. So when you need to evict someone to make room, you evict the minimum, which is always at the top. The K largest elements survive because they are never the minimum of the heap.
The complexity: you make n pushes and at most n pops, each O(log K). Total time: O(n log K). Compare to sorting everything: O(n log n). When K is much smaller than n, log K is much smaller than log n, and the heap is significantly faster. When K equals n, both are O(n log n) and there is no advantage.
This is the answer to “heap or sort”: use a heap when K is small relative to n and you need the top K of a large collection. Use sort when you need everything ordered, or when K is comparable to n.
Problem 1: Top K frequent elements
Find the K most frequent elements in an array. This is the two-step pattern from Episode 7’s drill: count frequencies with a hash map, then extract the top K by frequency using a heap.
from collections import Counter
import heapq
def top_k_frequent(nums, k):
counts = Counter(nums)
return [x for x, _ in heapq.nlargest(k, counts.items(), key=lambda p: p[1])]
top_k_frequent([1,1,1,2,2,3], 2) # -> [1, 2]
top_k_frequent([1,1,2,3,3,3], 1) # -> [3]
heapq.nlargest(k, iterable, key) implements exactly the size-K trick internally: it finds the K largest items by the given key without sorting the entire iterable. This is the idiomatic Python solution and runs in O(n log K) time.
The manual version makes the mechanics explicit:
def top_k_frequent_manual(nums, k):
counts = Counter(nums)
heap = []
for item, count in counts.items():
heapq.heappush(heap, (count, item)) # heap by count
if len(heap) > k:
heapq.heappop(heap) # evict least frequent
return [item for count, item in heap]
When you push tuples into a heap, Python compares the first element first (count), then the second (item) as a tiebreaker. The heap stays ordered by frequency. The size-K eviction keeps only the K most frequent.
The pattern. Hash map to count, heap to rank. This two-step combination solves “top K most/least/closest” problems that have a non-trivial ranking criterion. Count first, rank second. The heap does the ranking efficiently.
Problem 2: K closest points to origin
Given a list of points on a plane, find the K points closest to the origin. Distance is Euclidean: sqrt(x squared plus y squared). You do not need the square root for comparisons since sqrt is monotonic.
This is a min-distance top-K problem. You want the K with the smallest distance, which means a max-heap of size K (evict the furthest of the current top K when something closer arrives).
import heapq
def k_closest(points, k):
heap = []
for x, y in points:
dist_sq = x*x + y*y # squared distance, no sqrt needed
heapq.heappush(heap, (-dist_sq, x, y)) # negate: min-heap acts as max-heap
if len(heap) > k:
heapq.heappop(heap) # evict the farthest
return [[x, y] for _, x, y in heap]
k_closest([[1,3],[-2,2],[5,8],[0,1]], 2) # the 2 closest to origin
The negation trick: Python only provides a min-heap. To get max-heap behavior (evict the largest distance), negate the values so that the largest distance becomes the smallest negated distance and rises to the top. Evicting the minimum negated distance evicts the maximum actual distance, which is what you want.
This problem illustrates a key generalization: the top-K pattern is not only “largest” or “most frequent.” It is any problem where you want the K elements that rank highest by some criterion. Negate the criterion if you want the K smallest by some metric. The heap does not care about semantics, only about the comparison order.
A precision note from testing. When multiple points have equal distance to the origin, the heap may return them in a different order than sorting would. This is not a bug. The problem asks for any K valid points, not a specific tie-broken ordering. If a specific tie-breaking rule is required, the comparison needs to encode it explicitly.
Problem 3: Merge K sorted lists
You have K sorted linked lists. Merge them into one sorted linked list. This is the problem that directly extends Episode 6’s merge-two-sorted-lists (Tool 4) to K lists.
The naive extension merges lists pairwise: merge lists 1 and 2, merge the result with list 3, and so on. This is O(nK) where n is the total number of nodes: the final merge processes all n nodes, the second-to-last processes n minus n/K nodes, and so on.
The work grows linearly with K.
The heap extension: maintain a min-heap containing the current front node of each list. Always extract the global minimum front, attach it to the result, and push that list’s next node into the heap.
import heapq
class Node:
def __init__(self, val=0, nxt=None):
self.val = val
self.next = nxt
def __lt__(self, other):
return self.val < other.val # needed so heap can compare nodes
def merge_k_sorted(lists):
heap = []
for node in lists:
if node:
heapq.heappush(heap, node)
dummy = Node(0)
curr = dummy
while heap:
node = heapq.heappop(heap) # the globally smallest current front
curr.next = node
curr = curr.next
if node.next:
heapq.heappush(heap, node.next) # push the next node from this list
return dummy.next
The heap always contains at most K nodes (one per list). Each push and pop is O(log K).
Total work: n extractions and n insertions, each O(log K). Total time: O(n log K). For K lists of roughly equal size, this is a factor of K better than the naive pairwise approach.
The __lt__ method on Node is required because Python’s heap compares items with less-than, and when two nodes have the same value, it would try to compare the nodes themselves (which fails without this method).
Why K matters. The heap of size K is what makes this efficient. You never sort all n nodes. You only ever compare K front nodes against each other, and the heap does this in O(log K) per extraction. The pattern is the same as top-K largest: keep a small, efficient structure (heap of size K) rather than a large, slower structure (array of all K fronts compared linearly, which would be O(n times K)).
Problem 4: Find the median from a data stream
This is the hardest heap problem in common interviews. Numbers arrive one at a time in an arbitrary order, and after each arrival you must be able to report the current median in O(1). Insertions should be efficient.
The trick is two heaps: a max-heap for the lower half and a min-heap for the upper half.
The median is either the top of the max-heap (odd total), or the average of both tops (even total).
import heapq
class MedianFinder:
def __init__(self):
self.small = [] # max-heap (store negated) for the lower half
self.large = [] # min-heap for the upper half
def add(self, num):
heapq.heappush(self.small, -num) # push to lower half first
# ensure every element in small <= every element in large
if self.large and -self.small[0] > self.large[0]:
heapq.heappush(self.large, -heapq.heappop(self.small))
# balance sizes: small can have at most one more element than large
if len(self.small) > len(self.large) + 1:
heapq.heappush(self.large, -heapq.heappop(self.small))
elif len(self.large) > len(self.small):
heapq.heappush(self.small, -heapq.heappop(self.large))
def get_median(self):
if len(self.small) > len(self.large):
return float(-self.small[0])
return (-self.small[0] + self.large[0]) / 2.0
The invariant this maintains: every element in small is less than or equal to every element in large. Both heaps are kept balanced in size (differing by at most one). When the total count is odd, small has one extra element and its max (negated top) is the median. When even, the median is the average of both tops.
Two balancing steps after each insertion.
The first ensures the partition invariant: if the new element in small is larger than the smallest element in large, it belongs in large and we transfer it.
The second ensures the size invariant: if one heap grows too large, we transfer to the other.
Walk through inserting [5, 2, 4, 1, 3] one at a time:
Insert 5: small=[-5], large=[]. No large to compare. Size: 1-0, ok. Insert 2: small=[-5,-2] (push 2, negated=-2). large=[-(-5)]=[] wait, large is empty so skip partition check. Size: 2-0, unbalanced. Transfer top of small (5) to large. small=[-2], large=[5]. Median: 2+5/2 = 3.5. Correct (sorted: [2,5]). Insert 4: push -4 to small. small=[-4,-2]. -(-4)=4 > large[0]=5? No, 4<5, so partition ok. Size: 2-1, unbalanced. Transfer. small=[-2], large=[4,5]. Median: (2+4)/2=3. Correct (sorted: [2,4,5]). Insert 1: push -1. small=[-2,-1]. -(-2)=2 > 4? No, ok. Size 2-2, balanced. Median: (2+4)/2=3. Correct (sorted:[1,2,4,5]). Insert 3: push -3. small=[-3,-2,-1]. -(-3)=3 < 4, ok. Size 3-2, unbalanced. Transfer top (3) to large. small=[-2,-1], large=[3,4,5]. Median: -(-2)=2? No wait, len(small)=2, len(large)=3. Transfer top of large to small. small=[-3,-2,-1], large=[4,5]. Median: -(-3)=3. Correct (sorted:[1,2,3,4,5]).
After all insertions, get_median returns 3.0 (odd count, small has extra element, its negated top is 3). Verified against brute force on 200 random streams.
The two-heap technique handles a problem that a single heap cannot: finding the median requires knowing about the middle of the distribution, not the extreme.
By splitting the distribution at the median into two halves, each represented by its extreme (the max of the lower half and the min of the upper half), you can maintain the median in O(log n) insertions and O(1) queries.
A Precise Trace of the Size-K Trick
The size-K trick is described in one paragraph above, but tracing it precisely on a concrete example is worth the space because “use a min-heap of size K to find the K largest” is counterintuitive enough that many candidates understand it conceptually but implement it incorrectly under pressure.
Walk through top_k_largest([3, 1, 5, 2, 8, 7, 4], 3). K is 3.
Process 3: push 3. Heap: [3]. Size 1, no eviction needed. Process 1: push 1. Heap: [1, 3]. Size 2, no eviction needed. Process 5: push 5. Heap: [1, 3, 5]. Size 3, no eviction needed. Process 2: push 2. Heap: [1, 2, 5, 3]. Size 4, exceeds K. Evict minimum (1). Heap: [2, 3, 5]. Process 8: push 8. Heap: [2, 3, 5, 8]. Size 4, evict minimum (2). Heap: [3, 5, 8]. Process 7: push 7. Heap: [3, 5, 8, 7]. Size 4, evict minimum (3). Heap: [5, 7, 8]. Process 4: push 4. Heap: [4, 7, 8, 5]. Size 4, evict minimum (4). Heap: [5, 7, 8].
Final heap: [5, 7, 8]. These are the three largest. Correct.
Notice what happened at each eviction: the element evicted was always the smallest of the current top-K candidates, making room for the newcomer. Elements smaller than the current heap minimum (like 1, 2, 3, and finally 4) entered the heap and were immediately or soon evicted. Elements larger than the current minimum (8, 7) displaced smaller ones and stayed.
The invariant the heap maintains throughout: the heap contains the K largest elements seen so far. This invariant is preserved at each step because: if the new element is larger than the heap minimum, it replaces the minimum (which is now provably not in the top K); if the new element is smaller than or equal to the heap minimum, it is evicted immediately and the heap is unchanged.
This trace also reveals the correct implementation detail: you push first, then evict if size exceeds K.
The alternative (evict before pushing if the new element is smaller than the heap minimum) is tempting but wrong for equal values and edge cases. Always push first, then evict.
The Recognition Signal for Heap Problems
Four tells that indicate a heap:
“Top K,” “K most,” “K closest,” “K largest,” “K smallest,” “K most frequent.” The letter K next to a ranking word is the most reliable heap signal in all of interviewing. When you see it, your first question is whether K is small relative to the total (heap beats sort) or comparable (sort is equally good). In most interview problems, K is explicitly small.
“Repeatedly find the current minimum or maximum” from a changing collection. This is the stream variant. Something keeps arriving, and after each arrival you need to know the current extreme. A heap gives you O(log n) updates and O(1) extreme queries. Alternatives (resorted array, linear scan) are slower.
“Schedule tasks by priority” or “process next available event.” Any simulation or scheduling problem where you always process the highest-priority pending item is a heap. The heap is what gives you the next-highest-priority in O(log n) without re-scanning everything.
“Merge K sorted lists.” As shown above, any K-way merge where K sources each provide elements in sorted order is a heap problem. The heap keeps the K current fronts sorted efficiently.
Heap or Sort: the Explicit Comparison
The episode title is “Heap or Sort?”
Here is the direct comparison that answers it.
Use sorting when: you need all elements in order, K is comparable to n, the collection is static (fixed before processing), or the simplicity of sort matters more than the constant factor.
Use a heap when: you need only the top K elements from a large collection (K much smaller than n), elements arrive in a stream and you need continuous access to the current extreme, you need the K smallest or largest without a full sort, or you need the median of a stream.
The crossover is O(n log n) for sort versus O(n log K) for a size-K heap. When K equals n, both are the same. When K is ten and n is a million, sort does roughly twenty million operations and the heap does roughly seven million.
For K equal to a hundred and n equal to a billion, sort is infeasible and the heap processes it in about thirty billion operations (admittedly large, but linearly better per element).
In practice for interviews: if the problem gives you a specific small K and a large collection, the heap is the expected answer.
If K is not specified or could equal n, either approach is defensible and the simpler one (often sort) is preferred.
The Connection to Earlier Episodes
The heap deepens several connections the series has been building.
In Episode 6, merge two sorted lists used a two-pointer technique (Tool 4). Merge K sorted lists, covered here, is the K-way extension that requires a heap. The heap replaces the linear scan across K front nodes with a log-K lookup. This is the direct “Episode 6 Tool 4 scaled up” connection.
In Episode 4, BFS used a queue that explores nodes in order of discovery distance. Dijkstra’s shortest-path algorithm for weighted graphs replaces the BFS queue with a min-heap. The heap gives you the next-closest unvisited node (by accumulated weight) rather than the next node added. This is BFS upgraded: the heap provides the priority that BFS’s FIFO queue cannot.
In Episode 7’s drill, Problem 7 (top K frequent elements) previewed the hash map plus heap combination. This episode delivers it in full. The pattern, count first with a hash map, rank by count with a heap, is exactly what you apply there.
What Comes Next
Episode 9 covers backtracking: the “try, recurse, undo” pattern that generates all possibilities.
Backtracking is the complement of dynamic programming: DP counts or optimizes over possibilities without generating them, backtracking generates all of them explicitly.
The distinction between “how many ways” (DP, Episode 5) and “list all the ways” (backtracking, Episode 9) is the most important disambiguation in the series, and Episode 9 makes it concrete with subsets, permutations, and constraint-satisfaction problems.
Between Episodes 8 and 9, practice the heap recognition signal specifically. The top-K signal (K next to a ranking word) should trigger an automatic reach for the size-K heap. Sit down with one top-K problem you have not seen before, and before writing code, state: “this is a heap because of the K in top-K, and the approach is a size-K min-heap that evicts the smallest element when it grows beyond K.”
That narration is what pattern recognition looks like under interview conditions.
Frequently Asked Questions
Why does a min-heap find the K largest rather than the K smallest?
Because the min-heap of size K holds the K largest seen so far, with the smallest of those K at the top for easy eviction. When a new element arrives that is larger than the heap minimum, it displaces the minimum (the weakest of the current top K) and takes its place. If you want the K smallest, use a max-heap of size K (negate values to simulate it): the largest of the current bottom K is at the top for eviction when something smaller arrives.
What is heapq.nlargest and when should I use it?
It is Python’s built-in implementation of the size-K trick, finding the K largest items from an iterable by a given key. Use it when you have the full collection available at once. Use the manual loop when elements arrive in a stream and you need to maintain the heap incrementally.
Can I use a heap to sort a complete list?
Yes, this is heapsort. Push all elements, pop them one by one. O(n log n) time, O(1) space (in the in-place version). In practice, Python’s built-in sort (Timsort) is faster for complete sorting; use heapsort only when in-place O(1) space is specifically required.
In the MedianFinder, why push to small first?
Because the partition invariant (every element in small is ≤ every element in large) needs to be checked and possibly corrected after each insertion. Pushing to small first and then potentially moving the maximum to large is simpler to reason about than the reverse. Either direction works as long as the partition and size invariants are both restored after each insertion.
Does the heap guarantee that equal-distance points in k_closest are returned in a specific order?
No. When multiple points have equal distance, the heap uses the x and y coordinates as tiebreakers (for the tuple comparison), which may differ from other orderings. Problems that accept any K valid points are fine; problems that require a specific tie-breaking rule need to encode it explicitly in the comparison tuple.
When does Dijkstra’s algorithm use a heap and why?
Dijkstra’s finds shortest paths in a weighted graph. It always processes the unvisited node with the smallest accumulated distance so far. BFS does this trivially for unweighted graphs (every edge has the same cost, so the queue’s FIFO order corresponds to distance). For weighted graphs, FIFO order no longer corresponds to distance, so a min-heap by accumulated distance replaces the queue. Every BFS-to-Dijkstra conversion is a queue-to-heap swap.
The Bottom Line
A heap gives you O(1) access to the minimum (or maximum) element of a changing collection, with O(log n) insertions and deletions. It is the right tool whenever a problem needs the K best elements from a large collection (the size-K trick: a min-heap of size K holds the K largest, evicting the minimum when exceeded), whenever elements arrive in a stream and you need continuous access to the current extreme, or whenever K sorted sources need to be merged efficiently.
The core decision: heap when K is small relative to n (the size-K trick dominates), sort when you need everything ordered or K is comparable to n.
The recognition signal: K next to a ranking word (”top K,” “K closest,” “K most frequent”) means heap. “Repeatedly find the next minimum/maximum from a dynamic collection” means heap. “Merge K sorted sources” means heap.
Five problems in this episode: top-K largest (the canonical size-K example), top-K frequent (hash map plus heap), K closest points (the min-K variant with negation), merge K sorted lists (the Episode 6 extension), and median of a stream (two heaps, the hardest and most elegant application). All instances of the same underlying pattern: a small, efficient structure that keeps only what you need, rather than a large sort that arranges everything.
Episode 9 covers backtracking. Stay tuned.
Which of the five heap problems was hardest to see the heap in?