Patterns

Heaps and top-k

A heap is a tree stored so the smallest item, or the largest, can be removed in logarithmic time, and a new item can be inserted in logarithmic time. You do not use it to find an arbitrary element. You use it when you repeatedly need the current extreme value. Top-k problems are the interview shape of that idea: from a long stream, keep only the k items that are winning so far.

Python's heapq module is a min-heap. The smallest item sits at index 0. There is no max-heap type. You get the same effect by pushing the negation of a number, or by pushing tuples whose first element is the priority you want to order by. Say which one you are doing so the sign does not become a bug.

This page is not a graph algorithm page. Heaps show up inside shortest paths when edges have unequal weights, which the graphs page points at and then refuses to fake. Learn the container here. Attach it to a graph later if a prompt actually has weights.

Top-k with a heap of size k

Original task: a counter maps service names to how many requests they handled. Return the k busiest names, highest count first. If there is a tie, any order among the tied names is acceptable. If k is larger than the number of services, return all of them.

import heapq


def busiest(counts: dict[str, int], k: int) -> list[str]:
    if k <= 0:
        return []
    heap: list[tuple[int, str]] = []
    for name, count in counts.items():
        heapq.heappush(heap, (count, name))
        if len(heap) > k:
            heapq.heappop(heap)
    return [name for count, name in sorted(heap, reverse=True)]


assert busiest({"auth": 5, "feed": 9, "search": 7, "ads": 1}, 2) == ["feed", "search"]
assert busiest({"auth": 3}, 4) == ["auth"]
assert busiest({}, 3) == []

The heap never holds more than k services. Each of the n services is pushed once and maybe popped once. A push or pop is logarithmic in the heap size, which is at most k. Time is O(n log k). Extra space is O(k), plus the input dictionary which the caller already paid for. Sorting every service by count would be O(n log n). When k is much smaller than n, the heap is the better bound. When k is close to n, the sort is simpler and the bounds are similar. Say that trade. Choosing the sort when n is twenty is a fine interview answer if you say why.

The final sorted(heap, reverse=True) is a sort of k items, which is O(k log k) and does not change the overall bound. It exists because a heap is not fully sorted. Only the extreme element is in a known place. If the prompt only wanted the set of names, you could skip the sort. If it wants them in order, sort at the end or pop them off one by one. Popping from a min-heap yields the smallest count first, so you would reverse that list. The sort above is easier to read, and readability matters once the bound is already O(n log k).

Ties: (count, name) compares count first and the name only if counts match. That makes the heap's tie break deterministic, alphabetical by name, because strings compare lexicographically. The function then sorts again with reverse=True, which also orders names when counts match. The prompt said any order was acceptable, so this extra stability is a choice, not a requirement. Mention it if the interviewer cares about ties. Do not pretend the heap magically preserved input order.

The k-th largest value

A close cousin asks for a single rank rather than the whole list. The k-th largest element is the smallest item in a max-style set of k elements, which means a min-heap of size k still works. After the scan, the root of that heap is the k-th largest, because k items are at least that large and everything else was smaller and got popped.

def kth_largest(values: list[int], k: int) -> int:
    heap: list[int] = []
    for value in values:
        heapq.heappush(heap, value)
        if len(heap) > k:
            heapq.heappop(heap)
    if len(heap) < k:
        raise ValueError("k is larger than the input")
    return heap[0]


assert kth_largest([4, 1, 7, 3, 7], 2) == 7
assert kth_largest([4, 1, 7, 3], 1) == 7

Same complexity: O(n log k) time and O(k) space. The root is heap[0], not the result of heappop, if you still need the heap afterward. Popping would remove it. For a single return value, reading the root is enough.

The value error is a contract. Some prompts promise k is in range. Raising is still better than returning a wrong root from a short heap. An empty input with k equal to 1 should fail the contract. Try it.

Compare this with a full sort: sorted(values)[-k] is obvious and O(n log n). If you are under a clock and n is small, write the sort. The time management page would rather you finish. Switch to the heap when you need the tighter bound or when the values arrive as a stream and you cannot store them all. A stream is the setting where the heap is not just an optimization. You physically cannot sort what you did not keep.

What heaps are bad at

Heaps do not find the median in constant time just because you have heard of two-heap medians. That design is real, and it is also a specific structure: one heap for the lower half, one for the upper half, kept the same size. Do not start building it unless the prompt asks for a running median or something equivalent. A sorted list is allowed when the prompt is "sort these scores."

Heaps are also the wrong tool for "have I seen this key," which is a hash map, and the wrong tool for "the longest contiguous slice," which is a window. Top-k is about rank, not about position. If the prompt says contiguous, leave this page and open sliding window.

Priority queues in graph search have another trap: a node can be pushed more than once, with a better distance the second time. The simple heapq does not support decreasing a key in place. The usual Python approach is to push the new pair and ignore stale pairs when they come out. That is correct and a little wasteful. It is also more than this page needs to code. Knowing the sentence is enough until a prompt forces the issue.

Practice drill

Implement busiest without looking, then change the order to lowest count first and fix only the final sort. The heap condition stays. Then implement kth_largest with a full sort beside it and assert they match on a random list of a hundred numbers, for several k. Property tests catch sign errors that a single example can miss. Finally, explain in two sentences why the heap's space is O(k) and not O(n). If the explanation secretly stores every value, you have written a sort with extra steps.

The greedy page sometimes uses a heap when the next choice is "the interval that ends soonest" among many. You can solve the version on that page with a sort instead. Noticing that both exist is the skill. Reaching for heapq because the chapter title said heap is the common mistake.