Top-K Largest and Smallest

Every top-K question has the same shape: n values go in, the k best come out, and k is far smaller than n. Sorting answers it in O(n log n) time while holding all n values. The heap answer runs in O(n log k) time and holds k values. The idea behind it reads backwards at first: to keep the k largest, use a min-heap.

The root is the bar to clear

Cap a min-heap at k values. Its root is the smallest value being kept, so it is the weakest member of the current top-K, and a newcomer only deserves a place if it beats that member. A newcomer that wins evicts the root. A newcomer that loses cannot be in the top-K and is dropped after a single comparison. A max-heap would put the strongest kept value at the root, the one value the loop must never evict, and leave the weakest buried where no single comparison can reach it.

import heapq

def top_k(stream, k):
    heap = []
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:
            heapq.heapreplace(heap, x)
        print(x, heap)
    return sorted(heap, reverse=True)

result = top_k([4, 1, 7, 3, 8, 2, 9], 3)
print(result)
4 [4]
1 [1, 4]
7 [1, 4, 7]
3 [3, 4, 7]
8 [4, 8, 7]
2 [4, 8, 7]
9 [7, 8, 9]
[9, 8, 7]

Watch the root. After 4 and 1 arrive, then 7, the root is 1, and the value 3 evicts it. Next, 8 evicts the 3, the 2 loses to a root of 4, and 9 evicts the 4. The surviving heap [7, 8, 9] happens to be ascending, but heap order only promises the root, so the function sorts the survivors before returning [9, 8, 7].

What the heap saves

Sorting pays about log2 n comparisons for every value. The heap pays about log2 k for each value that gets in and a single comparison for each value that does not. In a pipeline the memory difference matters more: the heap needs k slots however long the stream runs, while a sort needs the entire input in memory at once. With n = 10,000,000 and k = 10:

The heap's worst case arrives when values come in ascending order, so every one evicts the root. On shuffled input, most values lose their first comparison and cost a single step each.

Smallest K is the mirror image

To keep the k smallest, the root must be the largest kept value, so cap a max-heap at k by storing negated values. A newcomer enters when its negation beats the root, which means the original is smaller than the largest value kept.

import heapq

def bottom_k(stream, k):
    heap = []
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, -x)
        elif -x > heap[0]:
            heapq.heapreplace(heap, -x)
    return sorted(-v for v in heap)

print(bottom_k([4, 1, 7, 3, 8, 2, 9], 3))

The library version and its edge cases

heapq.nlargest(k, iterable, key=None) and heapq.nsmallest(k, iterable, key=None) implement this pattern and return their results already sorted. The Python documentation adds 2 rules of thumb: max() and min() are more efficient when k = 1, and sorted() is more efficient for larger k. In an interview, write the loop to show the mechanism, then name nlargest as the call you would use in production.

Edge cases are where hand-written versions fail. With k = 0, the loop in the animation skips the push branch, compares the first value against heap[0] on an empty list and raises IndexError: list index out of range. With k larger than the input, the heap never fills and every value comes back. A guard brings the hand-written version in line with the library:

import heapq

def top_k(stream, k):
    if k <= 0:
        return []
    heap = []
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:
            heapq.heapreplace(heap, x)
    return sorted(heap, reverse=True)

data = [4, 1, 7, 3, 8, 2, 9]
print(top_k(data, 0), heapq.nlargest(0, data))
print(top_k(data, 20) == heapq.nlargest(20, data))

Ranking by a computed key

When the ranking value is computed, pass it as key instead of building tuples by hand. For the k points closest to the origin, the key is the squared distance, which ranks points in the same order as distance and skips the square root.

import heapq

points = [(3, 4), (1, 1), (-2, 0), (5, -1), (0, 2)]
closest = heapq.nsmallest(2, points, key=lambda p: p[0] ** 2 + p[1] ** 2)
print(closest)

Other ways interviewers phrase top-K

About this lesson

This section is part of the Heap & Top-K: 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.