Heap / Priority Queue

How it works

A heap keeps one extreme element (the smallest or the largest) available in O(1), while inserting and removing cost O(log n). It never sorts everything; it only maintains enough order to know what comes next. That is cheaper than sorting when you only need the next few items.

For top-k, the counterintuitive move is to keep a min-heap of size k. Each new item competes with the weakest of the current top k, which sits at the root. The heap never grows past k, so the cost is O(n log k) instead of O(n log n).

Heaps also merge many sorted sources. Seed the heap with the first element of each list; every time you pop the smallest, push the next element from the same list. The heap holds at most one candidate per source.

In disguise, a heap appears in scheduling (always run the task that finishes first), in simulations (repeatedly take the two largest stones), in running medians (two heaps balanced around the middle), and inside Dijkstra's algorithm.

How to attack it

1. 1Ask whether you repeatedly need the current minimum or maximum of a changing set.
2. 2Pick min-heap or max-heap, and the size bound if it is a top-k question.
3. 3Decide what to store: a value, or a tuple with a priority first and a tie-breaker.
4. 4Write the push and pop loop, and what triggers each.
5. 5Compare with sorting once: if the set never changes, a sort may be simpler.

Template

A skeleton to adapt, not a solution to paste. Write it from memory a few times.

import heapq


def k_largest(nums, k):
    # Min-heap of size k: the root is the weakest of the current top k.
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)    # drop the smallest
    return sorted(heap, reverse=True)


def merge_sorted(lists):
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            # (value, list index, position): the index breaks ties safely.
            heapq.heappush(heap, (lst[0], i, 0))
    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

# Python only has a min-heap: push -x to simulate a max-heap.

What it costs