Skip to content
BytePatterns

Top K Elements

Heaps: lesson 4 of 4

Hold k winners and let the weakest one guard the door.

Lesson 4 of 4 · 6 min

Top K Elements

Step 1 of 14

You rarely need everything ranked — just the best k = 3. Hold a min-heap of exactly three.

The Idea

You rarely need everything ranked — just the best k. Hold a min-heap of exactly k items, so the root is your weakest keeper.

Each new value is compared against that root alone. Bigger, and it evicts the weakling; smaller, and it is dropped for good. Cost: O(n log k) time and O(k) memory, with no need to hold the stream.

Real-World Example

A plant breeder walks a field of thousands of sunflowers carrying ten sample bags. Each promising head is weighed against the lightest bag already carried — the field never gets ranked, and the bags never outnumber ten.

The Code

import heapq

def top_k(stream, k):
    keep = []                                  # min-heap of the best k so far
    for x in stream:
        if len(keep) < k:
            heapq.heappush(keep, x)
        elif x > keep[0]:                      # beats the weakest keeper?
            heapq.heapreplace(keep, x)         # pop it, push x: one sift
    return sorted(keep, reverse=True)

print(top_k([4, 1, 9, 7, 3, 8], 3))            # [9, 8, 7]
print(heapq.nlargest(3, [4, 1, 9, 7, 3, 8]))   # same, batteries included

Your turn

What does this print?

import heapq
keep = [5, 6, 9]          # the 3 best so far
heapq.heapify(keep)
for x in [2, 7]:
  if x > keep[0]:
      heapq.heapreplace(keep, x)
print(keep[0])

Mini quiz

1 / 3

To keep the k largest values, which heap do you hold?