Heap vs Sort for Top-K

Concepts covered: Frequency Counting, Heap & Top-K, Lambda, Tuples

Prompts that say 'top 10', 'the 5 slowest', 'the 3 most frequent' or 'the highest scorer' ask for K rows, and ordering the other rows is wasted work. Sorting all n rows to read the first K does O(n log n) work and holds every row. A min-heap of size K does O(n log K) work and holds K rows, and Python's heapq module packages it as 2 functions.

import heapq

weekly = [
    {'track': 'Halo', 'streams': 30},
    {'track': 'Drift', 'streams': 50},
    {'track': 'Ember', 'streams': 45},
    {'track': 'Aster', 'streams': 30},
    {'track': 'Pulse', 'streams': 12},
]

top = heapq.nlargest(3, weekly, key=lambda r: r['streams'])
print([r['track'] for r in top])

low = heapq.nsmallest(2, weekly, key=lambda r: r['streams'])
print([r['track'] for r in low])

The first line prints ['Drift', 'Ember', 'Halo'] and the second prints ['Pulse', 'Halo']. Halo and Aster tie at 30 streams. The documentation defines nlargest(k, it, key=f) as equivalent to sorted(it, key=f, reverse=True)[:k], and nsmallest likewise without the reverse, so ties resolve as the stable sort would: Halo appears first in weekly, so Halo wins both times.

How the heap keeps only K rows

Keep a min-heap of the K largest values seen so far. Its smallest member sits at index 0 and acts as the bar to clear. Each new value is compared against that bar once. A value that fails is discarded after that single comparison; a value that beats it replaces the bar and sifts down through about log2 K levels.

import heapq

streams = [30, 50, 45, 30, 12, 61, 8]
k = 3
heap = []
for s in streams:
    if len(heap) < k:
        heapq.heappush(heap, s)
    elif s > heap[0]:
        heapq.heapreplace(heap, s)
    print(s, heap)
print(sorted(heap, reverse=True))

30 [30]
50 [30, 50]
45 [30, 50, 45]
30 [30, 50, 45]
12 [30, 50, 45]
61 [45, 50, 61]
8 [45, 50, 61]
[61, 50, 45]

Of the 4 values that arrived after the heap filled, only 61 forced a replacement; 30, 12 and 8 each cost 1 comparison. On real data most rows fail the first check, so the work in practice runs close to n comparisons.

The numbers behind the choice

For n = 10,000,000 and K = 10, a comparison sort needs about n log2 n, roughly 233 million comparisons. The heap's worst case is about n log2 K, roughly 33 million, and it holds 10 rows instead of a second 10,000,000-entry list. The advantage shrinks as K grows: the documentation notes that for larger K, sorted() is more efficient, and that for K = 1, min() and max() are the better calls.

• 7x
• 80 MB

Top K with a tie-break in the other direction

Prompt: the top 3 tracks by streams, ties broken by track name from A to Z. Calling nlargest with the key (streams, track) breaks ties from Z to A, because it maximizes the whole tuple, name included. Flip the problem: ask for the 3 smallest of (-streams, track). That is the negated tuple key from earlier in the lesson, handed to nsmallest.

import heapq

pairs = [('Halo', 30), ('Drift', 50), ('Ember', 45), ('Aster', 30), ('Pulse', 12)]

wrong = heapq.nlargest(3, pairs, key=lambda p: (p[1], p[0]))
right = heapq.nsmallest(3, pairs, key=lambda p: (-p[1], p[0]))
print([t for t, _ in wrong])
print([t for t, _ in right])

The first line prints ['Drift', 'Ember', 'Halo'], because (30, 'Halo') is larger than (30, 'Aster'). The second prints ['Drift', 'Ember', 'Aster'], the answer the prompt asked for.

Most frequent K

'Find the 3 most common error codes' is top K over counts. collections.Counter counts in one pass, and most_common(3) returns the 3 highest counts; given a number, CPython's implementation calls heapq.nlargest. The documentation guarantees that equal counts come back in the order their elements were first encountered.

from collections import Counter

codes = ['E42', 'E17', 'E42', 'E99', 'E17', 'E05', 'E42', 'E99']
print(Counter(codes).most_common(3))

This prints [('E42', 3), ('E17', 2), ('E99', 2)]. E17 and E99 tie at 2, and E17 comes first because it appeared first in codes.

Do

• Reach for nlargest or nsmallest when the prompt names a K
• Feed a generator to nlargest to keep memory at K rows
• Use nsmallest with a negated count for mixed-direction ties
• Use min() or max() when K is 1

Avoid

• Sort 10 million rows to read the first 10
• Assume nlargest breaks ties alphabetically
• Use a heap when the caller needs every row in order
• Forget that heapq only builds min-heaps

About this lesson

This section is part of the Custom Sorting: Beginner lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.