-
Notifications
You must be signed in to change notification settings - Fork 0
025 — Top K and Heap Pattern
The Heaps & Priority Queues chapter covered the data structure itself — how a heap keeps its min or max at the root in O(1) peek and O(log n) insert/remove. This chapter is about recognizing when to reach for one: any time you need the "best K" of something without fully sorting everything else.
Picture a talent show with 500 acts but only room to invite the top 5 to the finale. You don't rank all 500 acts from first to last — that's way more work than the question needs. Instead, you keep a small holding pen of exactly 5 finalists. As each new act performs, you compare it only against the weakest finalist currently in the pen. If the new act is better, the weakest one gets bumped out and the new act takes its place; otherwise you don't even bother remembering it. By the end, the pen holds exactly the top 5, and you never had to fully sort the other 495. That's a heap's superpower for this pattern: maintain a small window of "current best K" and let the root — the current weakest member of that window — be the only thing you ever compare against.
Reach for the top-K/heap pattern when you see:
- "K most frequent," "K least frequent," or "Kth largest/smallest element" — the question asks for a fixed-size subset or a single ranked element, not a full sort of everything.
- "Merge K sorted lists" or "merge K sorted arrays" — you repeatedly need "the smallest head among K candidates right now," which is exactly what a min-heap of size K gives you in O(log K) per step.
-
A running/streaming median or running top-K as data arrives one element at a time — "design a class that supports
addNumandfindMedian" signals a two-heap setup (a max-heap for the lower half, a min-heap for the upper half). - "Schedule tasks by priority," "closest K points to origin," or anything phrased as "top/closest/most urgent K out of N" — whenever N is large but K is small, and full sorting would be overkill (O(N log N)) compared to a heap-based approach (O(N log K)).
The unifying tell: you need the best (or worst) K items, or repeatedly need "the current extreme among a changing set," and K is meaningfully smaller than N — full sorting solves it too, but a heap solves it without paying for order you don't need.
Finding the 3 largest numbers in a stream using a min-heap of size 3 — the root is always the weakest of the current top 3, so a challenger only needs to beat the root:
stream: 5, 1, 9, 3, 7, 2, 8 keep top 3, min-heap capped at size 3
push 5: heap = [5] (root = 5)
push 1: heap = [1, 5] (root = 1)
push 9: heap = [1, 5, 9] (root = 1) <- full at size 3
push 3: 3 > root(1)? yes -> pop 1, push 3
heap = [3, 5, 9] (root = 3)
push 7: 7 > root(3)? yes -> pop 3, push 7
heap = [5, 7, 9] (root = 5)
push 2: 2 > root(5)? no -> discard 2
heap = [5, 7, 9] (unchanged)
push 8: 8 > root(5)? yes -> pop 5, push 8
heap = [7, 8, 9] (root = 7)
final top 3 = {7, 8, 9}
import Foundation
// A minimal min-heap wrapper — Swift's stdlib has no built-in heap,
// so problems typically reach for a sorted-insert Array (fine for small K)
// or roll a proper binary heap. Shown here: the "keep top K" driver logic,
// using an Array kept small and sorted as a stand-in for a heap of size K.
func topKTemplate<T>(_ items: [T], k: Int, isBetter: (T, T) -> Bool) -> [T] {
var window: [T] = [] // acts as a size-K min-heap; window[0] is the "weakest" kept item
for item in items {
if window.count < k {
window.append(item)
window.sort(by: isBetter) // 🔧 Fill in: swap for a real heap push for O(log k)
} else if isBetter(window[0], item) == false && isBetter(item, window[0]) {
// item beats the current weakest kept item — swap it in
window[0] = item
window.sort(by: isBetter) // 🔧 Fill in: swap for a real heap sift-down for O(log k)
}
}
return window // 🔧 Fill in: return top-K set, or just its extreme, per what's asked
}Top K Frequent Elements — given an integer array nums and an integer k, return the k most frequent elements.
import Foundation
func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
// 1. Count frequencies.
var freq: [Int: Int] = [:]
for n in nums {
freq[n, default: 0] += 1
}
// 2. Maintain a min-heap of size k, keyed by frequency.
// (Using a sorted array here for clarity; a real binary heap
// is a drop-in swap for the same driver logic.)
var heap: [(value: Int, count: Int)] = []
for (value, count) in freq {
if heap.count < k {
heap.append((value, count))
heap.sort { $0.count < $1.count } // heap[0] = weakest (lowest count) kept so far
} else if count > heap[0].count {
heap[0] = (value, count)
heap.sort { $0.count < $1.count }
}
}
return heap.map { $0.value }
}
// smoke test
print(Set(topKFrequent([1,1,1,2,2,3], 2)) == Set([1,2])) // true
print(Set(topKFrequent([1], 1)) == Set([1])) // true
print(Set(topKFrequent([4,1,-1,2,-1,2,3], 2)) == Set([-1,2])) // trueMapped onto the template: items is the list of (value, count) pairs from the frequency map, isBetter compares by count (higher frequency wins), and the size-k window keeps exactly the k most frequent values seen so far, discarding a challenger unless it beats the current weakest kept element. Since the final answer just needs the set of top-K values (order doesn't matter for this problem), the smoke test compares as Set rather than requiring an exact array order. For the full data-structure mechanics of the heap itself — array-backed layout, sift-up/sift-down — see Heaps & Priority Queues.
Top-K is the talent show finale pen — keep only the current best K, and let a new act beat only the weakest one in the pen, never the whole crowd.
For "Top K Frequent Elements" above, the window is a min-heap even though the problem asks for the most frequent elements. Explain why a min-heap (not a max-heap) is the right choice for finding the top K largest values, in terms of what needs to sit at the root so it can be evicted cheaply.
⬅️ Previous: Merge Intervals · Next: 1D Dynamic Programming ➡️