Coding Interview: Patterns & DSA · Sorting, Searching & Heaps · Lesson 2 of 3
Heaps & top-K
A heap (priority queue) always gives you the smallest (min-heap) or largest (max-heap) element in O(1), with O(log n) push and pop. It is the right tool whenever you repeatedly need the current extreme without fully sorting.
Top-K elements
To find the k largest elements, keep a min-heap of size k: push each element, and whenever the heap exceeds k, pop the smallest. What remains is the top k. This is O(n log k) — better than sorting the whole array (O(n log n)) when k is small.
import heapq def k_largest(nums, k): heap = [] # min-heap of size k for x in nums: heapq.heappush(heap, x) if len(heap) > k: heapq.heappop(heap) # drop the smallest return heap # the k largest values
- Top-K frequent / K closest: heap keyed by frequency or distance.
- Merge K sorted lists: a heap of the current heads, pop the smallest, push its successor.
- Median of a stream: two heaps (a max-heap of the low half, a min-heap of the high half) balanced in size.
Why not just sort?Sorting to grab the top
k is O(n log n) and throws away most of the work. A size-k heap is O(n log k) and uses O(k) space — a clear win when k << n. (Quickselect gets top-K in O(n) average if you do not need them sorted.)Recognize itHeaps are the tell for top-K, K closest, merge K sorted, running median, and scheduling by priority — anything repeatedly asking for the current min or max.
◆ Lock it in
- A heap gives the running min/max in
O(1), push/pop inO(log n). - Top-K with a size-k heap is
O(n log k)— beats sorting whenkis small. - Two balanced heaps track a streaming median; a heap merges K sorted lists.
Feynman drill — say it out loudExplain why a size-k min-heap finds the k largest elements in O(n log k) rather than O(n log n).
Step 1 rate your confidence · Step 2 pick your answer
To keep the k LARGEST elements of a stream, you maintain a size-k:
Step 1 — how sure are you?