Top K Elements With a Heap: Why Not Just Sort?

"Return the k largest" has an answer that everyone writes first: sort, slice, done. It is correct, it is one line, and in an interview it invites the follow-up "can you do better?" The answer is yes, and the reason is worth understanding rather than memorising: sorting answers a much bigger question than the one you were asked.

The problem it solves

Given n items, return the k largest — the ten hottest products, the hundred closest points, the five most frequent words. Usually k is tiny next to n, and often the items arrive as a stream you cannot hold in memory all at once.

Sorting produces the complete ranking of all n items, in O(n log n). Then you throw away all but k of them. You have paid to decide whether item 50,000 beats item 50,001, and nobody asked.

The intuition

Carry only the current best k. For each new item the only question is: does it belong among them? And to answer that you only need to compare against the weakest of the current keepers. If the newcomer cannot beat the weakest, it cannot beat any of them.

So the structure you want gives instant access to the smallest of your k keepers and lets you replace it cheaply. That is a min-heap of size k — which surprises people, because the question is about the largest.

The heap's root is the bouncer. It is the lowest bar a newcomer must clear, and it is the one who leaves when someone better arrives.

Most items fail at the door with one comparison and never touch the heap. Only an item that beats the root triggers a replacement, and a replacement costs O(log k) — the heap is only k tall.

Watch it run

Watch the stream rather than the heap. Most values are compared once against the root and dropped; a few get in and push the weakest keeper out. The heap never grows past k, however long the stream.

The code

The heap solution, checked against the sort, with every comparison counted.

import heapq
from collections import Counter

def top_k(stream, k):
    keep = []                                  # min-heap: weakest keeper on top
    for x in stream:
        if len(keep) < k:
            heapq.heappush(keep, x)
        elif x > keep[0]:                      # beats the weakest keeper?
            heapq.heapreplace(keep, x)         # evict it, sift x into place
    return sorted(keep, reverse=True)

class Counted:                                 # a number that counts comparisons
    calls = 0
    def __init__(self, v): self.v = v
    def __lt__(self, other):
        Counted.calls += 1
        return self.v < other.v
    def __gt__(self, other):
        Counted.calls += 1
        return self.v > other.v

n, k = 100_000, 10
data = [(i * 7919) % n for i in range(n)]      # a fixed shuffle of 0..n-1

Counted.calls = 0
by_sort = [c.v for c in sorted(map(Counted, data), reverse=True)[:k]]
sort_cost = Counted.calls

Counted.calls = 0
by_heap = [c.v for c in top_k(map(Counted, data), k)]
heap_cost = Counted.calls

print(by_sort == by_heap, by_heap[:3])         # True [99999, 99998, 99997]
print(sort_cost, heap_cost)                    # 1515712 100431

Counted.calls = 0
top_k(map(Counted, range(n)), k)               # ascending: every value evicts
print(Counted.calls)                           # 439974

words = "the cat and the hat and the bat".split()
print(Counter(words).most_common(2))           # [('the', 3), ('and', 2)]

Same answer, fifteen times fewer comparisons. On the shuffled input almost every value is rejected by the single x > keep[0] check — about 100,000 comparisons for 100,000 values, plus a few hundred for the rare replacements.

The third number is the heap's worst case. On ascending input every new value beats the root, so every one triggers a full replacement. Even then it is about 440,000 comparisons — n log k rather than n log n, still under a third of the sort.

Two practical notes. heapq.heapreplace pops and pushes in one sift, which is cheaper than calling heappop and then heappush. And the standard library already packages this exact idea as heapq.nlargest(k, iterable), which is what you should use in real code once you have shown you know what it does.

The complexity

• Sort and slice: O(n log n) time, O(n) memory — the whole input must be materialised.
• Size-k min-heap: O(n log k) time, O(k) memory. It works on a stream of unknown length, because it never needs more than the current k.
• Quickselect: O(n) on average, by partitioning around a pivot until the k-th position is found. Faster in theory, but it needs the whole array in memory, rearranges it, and has an O(n²) worst case unless the pivot is chosen carefully.
• Buckets, for top-k frequent. Counts are bounded by n, so grouping items by count and reading buckets from the top is O(n) with no heap at all — the bucket version of the question.

When k approaches n, log k approaches log n and the heap's advantage disappears. If you need the top half, just sort.

Where it goes wrong

• Using a max-heap of everything. Heapify all n, pop k times: O(n + k log n). That is also fine on time — but it holds all n items, which defeats the streaming case.
• A max-heap of size k. Its root is the strongest keeper, which is the wrong one to compare against. The heap direction is opposite to the question: min-heap for the largest k, max-heap for the smallest k. Python's heapq only has a min-heap, so the smallest-k version pushes negated values.
• Forgetting the result is not sorted. A heap is only partially ordered. If the caller wants the k items in rank order, sort the k at the end — O(k log k), negligible.
• Ties. "Top k frequent words" usually specifies how to break ties, often alphabetically. Push (count, word) tuples and check that the tuple ordering matches the rule — for a min-heap of the largest counts, alphabetical tie-breaking needs the word compared in reverse, which is the classic trap.

How to say it in an interview

"Sorting works but it ranks all n items when I only need k. Instead I keep a min-heap of size k: its root is the weakest of my current top k. Each new item is compared with the root; if it is larger, it replaces the root in O(log k), otherwise it is discarded. That is O(n log k) time and O(k) space, and it works on a stream. If memory were not a concern and I needed only the k-th element, quickselect would be O(n) on average."

Then add the one sentence that shows judgement: when k is close to n, sorting is simpler and just as fast.