The Top-K Elements Pattern

When you only need the biggest, smallest, or most frequent handful of items out of a much larger collection, sorting the whole thing is overkill. The Top-K Elements pattern solves this with a heap that never grows past size k, answering the question in O(n log k) time instead of the O(n log n) a full sort would cost. It shows up everywhere: trending hashtags, the k nearest stores to a user, the k most frequent words in a document, or the top-scoring candidates in a ranking system. This lesson builds the pattern from first principles and applies it to three realistic problems.

Overview: How the Top-K Pattern Works

Suppose you have a list of a million numbers and need the 10 largest. Sorting the entire list costs O(n log n) and gives you far more information than you need — a total order over a million numbers when you only wanted 10 of them. The Top-K pattern instead keeps a small heap that never holds more than k elements, and it relies on a clever inversion: to track the k largest values seen so far, you use a min-heap of size k.

Here’s the intuition. Walk through the input one element at a time, pushing each onto a min-heap. Whenever the heap’s size exceeds k, pop the smallest element off. Because you always evict the smallest item currently in the heap, only the largest k items you’ve seen survive — and the heap’s root (the minimum of that group) is the cutoff: any future element smaller than the root can never make the top k, so it gets pushed and immediately popped right back out. Wanting the k smallest values instead just flips the logic: use a max-heap of size k (in Python, simulate one by pushing negated values, since heapq only implements a min-heap) and evict the largest element whenever the heap overflows.

This “heap capped at size k” idea generalizes beyond raw largest/smallest. If you attach a priority to each item — a frequency count, a distance, a score — the same eviction logic finds the top k by that priority. That’s why this pattern pairs naturally with collections.Counter for frequency-based problems, and why it’s one of the most common building blocks in coding interviews.

heapq.nlargest and heapq.nsmallest

Python’s standard library already implements this pattern. heapq.nlargest(k, iterable) and heapq.nsmallest(k, iterable) return the k largest or smallest items from any iterable, optionally with a key function for custom priorities. When k is small relative to the input, they use exactly the capped-heap technique described above; when k is close to the size of the whole input, they fall back to sorting everything, since sorting is no less efficient at that point. Reach for these functions for one-off queries — hand-roll the heap yourself when you need a running top-k over a stream, where you can’t just call a function once and be done.

Time and Space Complexity

Let n be the number of elements you scan and k be how many results you want to keep (usually k is much smaller than n).

Operation Time Space Why
Build a size-k heap over n elements O(n log k) O(k) Each of the n elements triggers at most one push and one pop, each costing O(log k) on a heap that never exceeds size k.
Single push/evict on a running stream O(log k) O(1) amortized per call The heap already holds at most k items, so any single insertion or removal walks at most log2(k) levels.
heapq.nlargest(k, n items) O(n log k) O(k) Same capped-heap mechanics as the hand-rolled version, implemented in C for speed.
Sort everything, then slice the top k O(n log n) O(n) Timsort (Python’s sort) must fully order all n elements even though only k are used afterward.

The gap between O(n log k) and O(n log n) matters more than it looks: if n is a million and k is 10, log k is about 3.3 while log n is about 20 — roughly a 6x difference in comparisons, and the heap only ever holds 10 items instead of a million. That’s also why the heap approach uses O(k) space instead of O(n): the whole input never needs to be materialized or fully sorted at once, which matters when the input is a live stream too large to hold in a sorted structure.

Examples

Example 1: The k Largest Elements With a Capped Min-Heap

This is the pattern in its purest form: push every element onto a min-heap, and whenever the heap grows past size k, pop the smallest item.

import heapq

def top_k_largest(nums: list[int], k: int) -> list[int]:
    heap: list[int] = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return sorted(heap, reverse=True)

nums = [3, 1, 5, 12, 2, 11]
result = top_k_largest(nums, 3)
print(result)

Output:

[12, 11, 5]

Trace it: the heap fills with 3, 1, 5 after the first three pushes. Pushing 12 grows it to size 4, so the smallest value, 1, is popped, leaving {3, 5, 12}. Pushing 2 grows it again; 2 is now the smallest, so it’s popped right back out, leaving {3, 5, 12} unchanged. Pushing 11 grows it once more; this time 3 is the smallest and gets evicted, leaving {5, 11, 12}. Those three survivors are exactly the three largest values in the input, and sorting them in reverse for display gives [12, 11, 5].

Example 2: Top k Frequent Elements

A very common interview variant ranks by how often each value appears rather than by its value. Count occurrences with Counter, then ask heapq.nlargest for the k keys with the highest counts.

from collections import Counter
import heapq

def top_k_frequent(nums: list[int], k: int) -> list[int]:
    counts = Counter(nums)
    return heapq.nlargest(k, counts.keys(), key=counts.get)

nums = [1, 1, 1, 2, 2, 3, 4, 4, 4, 4]
result = top_k_frequent(nums, 2)
print(result)

Output:

[4, 1]

Counter(nums) tallies each value’s frequency: 1 appears 3 times, 2 appears 2 times, 3 appears once, and 4 appears 4 times. heapq.nlargest(2, counts.keys(), key=counts.get) asks for the 2 keys whose count is largest — that’s 4 (count 4) and 1 (count 3), giving [4, 1].

Example 3: Kth Largest Element in a Stream

Streaming problems are where hand-rolling a heap earns its keep, because you can’t just call nlargest once — new values keep arriving and you need an answer after every one. This class keeps a min-heap capped at size k; its root is always the current kth largest value seen so far.

import heapq

class KthLargest:
    def __init__(self, k: int, nums: list[int]) -> None:
        self.k = k
        self.heap: list[int] = []
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

kth_largest = KthLargest(3, [4, 5, 8, 2])
print(kth_largest.add(3))
print(kth_largest.add(5))
print(kth_largest.add(10))
print(kth_largest.add(9))
print(kth_largest.add(4))

Output:

4
5
5
8
8

The constructor seeds the heap with [4, 5, 8, 2] one value at a time through add, so after construction the heap holds the 3 largest of those four: {4, 5, 8} (2 was pushed and immediately evicted as the smallest of four). Calling add(3) pushes 3, the heap now has four values {3, 4, 5, 8}, and the smallest, 3, is popped, leaving {4, 5, 8} — root 4. add(5) pushes another 5, evicts the smallest 4, leaving {5, 5, 8} — root 5. add(10) pushes 10, evicts 5, leaving {5, 8, 10} — root 5. add(9) pushes 9, evicts 5, leaving {8, 9, 10} — root 8. add(4) pushes 4, evicts it right back out since it’s now the smallest, leaving {8, 9, 10} unchanged — root 8.

How It Works, Step by Step

Walk through top_k_largest([3, 1, 5, 12, 2, 11], k=3) from Example 1 one push at a time, tracking the heap’s contents as a set (the exact internal array order heapq uses is an implementation detail — what matters is which elements survive):

Step Action Heap contents after step
1 push 3 {3}
2 push 1 {1, 3}
3 push 5 {1, 3, 5}
4 push 12, size 4 > k, pop min (1) {3, 5, 12}
5 push 2, size 4 > k, pop min (2) {3, 5, 12}
6 push 11, size 4 > k, pop min (3) {5, 11, 12}

Notice step 5: pushing 2 briefly makes the heap size 4, but 2 is immediately the smallest element and gets popped back out without ever displacing a real member of the top-3 group. That’s the pattern doing its job — values that can never make the cut are discarded in O(log k) time without disturbing the survivors. After all six elements are processed, {5, 11, 12} is exactly the set of the three largest values in the input.

Common Mistakes

Mistake 1: Off-by-one in the size check

It’s tempting to evict as soon as the heap reaches size k instead of after it exceeds k. That single-character bug throws away one real element every time and leaves the heap one short of what was asked for.

import heapq

def top_k_largest_buggy(nums: list[int], k: int) -> list[int]:
    heap: list[int] = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) >= k:
            heapq.heappop(heap)
    return sorted(heap, reverse=True)

nums = [3, 1, 5, 12, 2, 11]
print(top_k_largest_buggy(nums, 3))

Output:

[12, 11]

Using >= k means an eviction fires the moment the heap has only k elements — before a (k+1)th element has even arrived to compete for a spot. The heap reaches size 3 on the push of 5 and immediately evicts back down to size 2 on every push after that, so it never actually holds 3 candidates at once. The result has only 2 elements instead of the requested 3. The fix is to evict only once the heap has grown past k:

import heapq

def top_k_largest_fixed(nums: list[int], k: int) -> list[int]:
    heap: list[int] = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return sorted(heap, reverse=True)

nums = [3, 1, 5, 12, 2, 11]
print(top_k_largest_fixed(nums, 3))

Output:

[12, 11, 5]

Mistake 2: Assuming heapq can be configured as a max-heap

Python’s heapq module implements a min-heap only — there’s no argument or flag to flip it. Code that pops from a freshly-heapified list expecting the largest values back-to-back actually gets the smallest ones.

import heapq

def top_3_largest_wrong(nums: list[int]) -> list[int]:
    heap = list(nums)
    heapq.heapify(heap)
    result = []
    for _ in range(3):
        result.append(heapq.heappop(heap))
    return result

nums = [3, 1, 5, 12, 2, 11]
print(top_3_largest_wrong(nums))

Output:

[1, 2, 3]

heapq.heapify arranges the list so heapq.heappop always returns the current minimum — that’s true regardless of what the values represent, so three pops in a row return the 3 smallest values (1, 2, and 3), the exact opposite of what was wanted. The standard fix is to negate every value before heapifying, turning “smallest negated value” into “largest original value”:

import heapq

def top_3_largest_fixed(nums: list[int]) -> list[int]:
    negated = [-num for num in nums]
    heapq.heapify(negated)
    result = []
    for _ in range(3):
        result.append(-heapq.heappop(negated))
    return result

nums = [3, 1, 5, 12, 2, 11]
print(top_3_largest_fixed(nums))

Output:

[12, 11, 5]

Negation works cleanly for plain numbers; for more complex priorities (tuples, objects), it’s often clearer to just use heapq.nlargest instead of hand-managing a negated heap.

Best Practices

  • Cap the heap at size k, not n. The whole point of this pattern is O(k) space and O(n log k) time — building a heap over the full input and popping k times gives up the space savings and often the time savings too.
  • Use a min-heap to track the k largest values and a max-heap (negate values, or use heapq.nlargest) to track the k smallest — it’s easy to get this backwards, so double-check the direction before writing the eviction condition.
  • For a single, one-shot query, prefer heapq.nlargest(k, iterable, key=...) or heapq.nsmallest(k, iterable, key=...) over hand-rolling the loop — they’re implemented in C and handle the “k close to n, might as well sort” fallback for you.
  • For a running/streaming top-k where new values keep arriving, keep a persistent heap object across calls instead of recomputing from scratch — that turns an O(n log n) resort per update into an O(log k) update.
  • When ranking by a derived priority (frequency, distance, custom score) rather than the raw value, push tuples like (priority, item) or use the key argument — don’t compare raw items directly if they aren’t naturally orderable.
  • If two items can share a priority and you push plain (priority, item) tuples where item isn’t comparable (custom objects, dictionaries), Python raises a TypeError the moment it needs to break the tie. Add an explicit tiebreaker field, such as an insertion index.
  • Handle the edge cases: k = 0 should return an empty result without touching the heap, and k >= n means every element belongs in the “top k” — no eviction ever needs to happen.

Practice Exercises

  1. Write top_k_smallest(nums: list[int], k: int) -> list[int] that returns the k smallest values from a list, using a heap that never exceeds size k. Hint: you’ll want to evict the largest value whenever the heap overflows — the mirror image of Example 1.
  2. Given a list of words, write top_k_frequent_words(words: list[str], k: int) -> list[str] that returns the k most frequent words, breaking ties by choosing the alphabetically smaller word first. Hint: design the sort key so both frequency and the tiebreak compare in the direction you want when used with heapq.nlargest‘s key argument.
  3. Design a class StreamKthSmallest that supports add(val: int) -> int, returning the current kth smallest value seen so far across all calls (assume at least k values have been added before you rely on the return value). Hint: this is the mirror image of the KthLargest stream class in Example 3 — what kind of heap tracks the k smallest values as a running set?

Summary

  • The Top-K pattern keeps a heap capped at size k instead of sorting the whole input, cutting cost from O(n log n) to O(n log k) time and O(n) to O(k) space.
  • To find the k largest values, use a min-heap of size k and evict the smallest element whenever the heap overflows; for the k smallest, flip to a max-heap (negate values in Python) and evict the largest.
  • Python’s heapq.nlargest and heapq.nsmallest already implement this pattern for one-shot queries; hand-roll a heap only when you need a running top-k over a stream.
  • Combine with collections.Counter to rank by frequency instead of raw value — a very common interview shape.
  • Watch for the off-by-one eviction bug (>= k vs > k) and remember heapq is min-heap only — there is no built-in max-heap mode.