Top-K Problems — heap, sort, or quickselect?
Top-K Problems
"Give me the k largest values" — the most common interview and production question. Three approaches; pick by how k compares to n and whether data is streaming.
The three approaches
heapq.nlargest(k, data) — O(n log k). Maintains min-heap of size k. Memory O(k). Streaming winner: no full input in memory, and when k << n, log k demolishes log n.
sorted(data)[-k:] — O(n log n). Simple, Timsort constants excellent. When k ~ n, sort wins on hardware despite worse asymptotics.
Quickselect — O(n) average, O(n²) worst. Partitions around pivot, recurses into target side. Mutates input, doesn't preserve order. Random pivot mandatory. Choose when you need the k-th element specifically.
How to choose
- k << n, streaming, huge n:
heapq.nlargest. - k ~ n, or need them sorted:
sorted. - Only the k-th element (single value):
quickselect. - Top-k by key:
heapq.nlargest(k, data, key=fn).
Common bugs
- Sorting when you should heap.
sorted(data)[-k:]on 100M rows loads everything. - Off-by-one on "k-th largest". In length-n array, k-th largest = (n-k)-th smallest.
- Deterministic pivot.
arr[lo]blows up on sorted input. - Mutation surprise. Quickselect rearranges input. Pass
.copy().
Try it
- Time all three on n=200_000 with k=10 and k=150_000.
- Find k most frequent in a list of a million ints.
collections.Counter+heapq.nlargest(k, counter.items(), key=lambda kv: kv[1]).
3 graded exercises
Write the code, click Check. We’ll tell you exactly what to fix.
- Exercise 1
Write a function
k_largest(nums, k)that returns the k largest numbers from the listnums, sorted in descending order. Useheapq.nlargest. - Exercise 2
Write
top_k_frequent(words, k)that returns the k most-frequently occurring strings inwords. Ties may be broken in any order; return a list of length k. - Exercise 3
Write
kth_largest(nums, k)that returns the kth largest element innums(1-indexed — k=1 is the maximum). Use a heap of size k for O(n log k) time.