Heap patterns: top-k and streams

Top k: why a heap of size k

Sorting all n values and taking the first k costs O(n log n) and needs all data in memory. When k is small, there is a better way. Keep a min-heap with at most k items. For each new value, push it; if the heap now has k + 1 items, pop the smallest. At the end the heap holds exactly the k largest values, and its root heap[0] is the k-th largest.

Why a MIN-heap for the LARGEST values? The root is the smallest of the current top k, the first one to be replaced when a bigger value arrives. Each push and pop is O(log k), so the total is O(n log k) time and O(k) space. That works on a stream too: you never need more than k values in memory.

import heapq

def k_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)        # drop the smallest of the k + 1
    return sorted(heap, reverse=True)

print(k_largest([3, 2, 1, 5, 6, 4], 2))   # Output: [6, 5]

K closest points and other "best k by score" problems

When items have a score, compare by the score. For the k points closest to the origin, the score is x² + y² (no square root needed, because squaring keeps the order). heapq.nsmallest(k, points, key=...) is the quick answer. Written by hand, keep a MAX-heap of size k on distance (push the negative distance) and pop the farthest whenever the heap grows beyond k.

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

Top k frequent items combines Module 3 and this topic: count with collections.Counter in O(n), then pick the k items with the highest counts with a heap. Always say how ties are broken, for example "equal counts: smaller value first", and make your key match it.

Merge k sorted lists

You have k sorted lists with N values in total, for example sorted logs from k servers. Put the FIRST item of every list into a min-heap as (value, list number, position). Pop the smallest, add it to the output, and push the next item from the same list. The heap never holds more than k items, so each step is O(log k): O(N log k) in total. The list number in the tuple also breaks ties, so equal values never fall back to comparing anything strange.

def merge_k(lists):
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)
    out = []
    while heap:
        val, i, j = heapq.heappop(heap)
        out.append(val)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return out

print(merge_k([[1, 4, 5], [1, 3, 4], [2, 6]]))   # Output: [1, 1, 2, 3, 4, 4, 5, 6]

Median of a data stream: two heaps

The median is the middle value of the sorted data (the average of the two middle values when the count is even). Re-sorting after every new value costs O(n log n) each time. Instead keep two heaps: low, a max-heap (stored negated) with the smaller half, and high, a min-heap with the larger half. Keep their sizes equal, or let low hold one extra. Then the median is the top of low, or the average of both tops.

To add a number: push it into low, move low's largest into high (so every value in low stays ≤ every value in high), and if high is now bigger than low, move high's smallest back. Each add is O(log n) and reading the median is O(1).

class MedianFinder:
    def __init__(self):
        self.low = []    # max-heap of the smaller half (negated)
        self.high = []   # min-heap of the larger half

    def add_num(self, num):
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def find_median(self):
        if len(self.low) > len(self.high):
            return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2

m = MedianFinder()
for x in [5, 15, 1, 3]:
    m.add_num(x)
    print(m.find_median(), end=" ")   # Output: 5.0 10.0 5.0 4.0

Task scheduling

Schedulers often keep jobs in a heap keyed by time or priority: the next job to run is always at the root. A common interview version gives tasks with counts and a cooldown between two runs of the same task, and asks for the shortest schedule. The greedy plan is to always run the remaining task with the HIGHEST count that is not cooling down, using a max-heap of counts and a queue of tasks waiting out their cooldown.

Common mistakes

• Using a max-heap of all n items to find the k largest. It works but costs O(n + k log n) memory and time on a stream. Fix: a min-heap of size k.
• Returning the heap list as the answer when the order matters. The heap is not sorted. Fix: sort the k items, or pop them in order.
• Leaving ties unstated in "top k frequent". Fix: state the rule and use a key like (-count, value).
• In merge k, pushing (value, list) tuples: equal values then compare the lists themselves. Fix: include the list index.
• In the two-heap median, letting the heaps get out of balance by more than one, or forgetting to negate when reading low. Fix: rebalance after every add, and test with 1, 2 and 3 values.