Top K with a Heap

Top K with a heap is keeping only the K largest values seen so far. The heap holds the current winners, and smaller values are dropped as they arrive.

Recognition signal

A problem fits this pattern when it asks for the K largest or K smallest items. The input may be long, but only K answers are needed at the end.

signals = [
    "Asks for the K largest items.",
    "Input is long, answers are few.",
    "Only K winners are kept."
]
for line in signals:
    print(line)

Template: keep K largest

The template pushes each value onto a heap and pops the smallest when the heap grows past K. Keeping K winners is the whole pattern in plain words.

import heapq

def top_k(nums, k):
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return sorted(heap)

print(top_k([5, 1, 9, 3, 7, 2], 3))

Worked problem: top 2 scores

The same template works on a short list of scores. Only the two largest survive, and the rest are dropped one by one.

import heapq

def top_k(nums, k):
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return sorted(heap)

print(top_k([10, 4, 22, 7, 15], 2))

Example

import heapq

words = ["red", "blue", "red", "green", "blue", "red"]
counts = {}
for word in words:
    counts[word] = counts.get(word, 0) + 1

heap = []
for word, count in counts.items():
    heapq.heappush(heap, (count, word))
    if len(heap) > 2:
        heapq.heappop(heap)

result = sorted(heap, key=lambda item: (-item[0], item[1]))
print([(word, count) for count, word in result])

Key points

Heap: a structure that keeps its smallest item on top.
heapq: the Python module that turns a plain list into a heap.
Use a heap of size K to track the K largest values in a stream.
Push each new value, then drop the smallest item when the heap grows past K.

Counts are red 3, blue 2, and green 1. The heap keeps only two pairs, so green is dropped as the smallest. Sorting the survivors by count gives red first and blue second.

Note

at the end, the heap holds exactly the K largest values, unsorted.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import heapq

words = ["red", "blue", "red", "green", "blue", "red"]
counts = {}
for word in words:
    counts[word] = counts.get(word, 0) + 1

heap = []
for word, count in counts.items():
    heapq.heappush(heap, (count, word))
    if len(heap) > 2:
        heapq.heappop(heap)

result = sorted(heap, key=lambda item: (-item[0], item[1]))
print([(word, count) for count, word in result])

Exercise

A stream has 100 numbers, and you want the 3 largest. How large should the heap stay?

3 100 1 0