Heaps and heapq

Some programs need the smallest item again and again while new items keep arriving: a printer that prints the shortest job next, a game that runs the next event in time order, an emergency room that treats the most urgent patient first. A collection like that is a priority queue. With a plain list, min finds the smallest in O(n) every time. A list kept sorted does no better in Big-O terms: each new item costs O(n), whether bisect.insort puts it in its place (every item after it moves, which is quick work in C but still O(n)) or you append it and sort again (Timsort sees that the list is almost in order, but still looks at every item). A heap adds an item in O(log n), removes the smallest in O(log n), and shows the smallest in O(1).

A tree in a list

A binary heap is a binary tree with two rules:

• Shape: every level is full, except maybe the last, which fills from the left. So the height is about log₂ n.
• Order: every parent is smaller than or equal to its children. So the smallest value is at the root. (Unlike in a BST, nothing orders the left side against the right.)

With no gaps in the tree, it fits in a plain list, level by level, without any Node objects: the root at index 0, the children of index i at 2 * i + 1 and 2 * i + 2, and its parent at (i - 1) // 2.

index:  0  1  2  3  4
       [1, 3, 8, 5, 4]

           1
         /   \
        3     8
       / \
      5   4

To push a value, append it (the next free place on the last level) and sift it up: while it's smaller than its parent, swap the two. To pop the smallest, take the root, move the last item into its place and sift it down: while it's greater than the smaller of its children, swap it with that child. Each walks one path between the root and the bottom: O(log n).

Why the smaller child? After the swap, that child is the parent of its old sibling, and it is no greater than that sibling, so the order rule holds there. In the heap above, popping 1 moves 4 to the root, above 3 and 8. Swapping 4 with 3 gives [3, 4, 8, 5], a heap again. Swapping it with 8, the larger child, would give [8, 3, 4, 5], with 8 above 3. Nothing stops with an error, but the next pop returns 8 instead of 3. Always swapping with the left child goes wrong in the same way whenever the right child is the smaller one. Sifting up keeps the rule for a similar reason: the value that moves up is smaller than its old parent, which was no greater than its other child.

def push(heap, value):
    heap.append(value)
    i = len(heap) - 1
    while i > 0:
        parent = (i - 1) // 2
        if heap[i] >= heap[parent]:
            break
        heap[i], heap[parent] = heap[parent], heap[i]
        i = parent


def pop(heap):
    smallest = heap[0]
    last = heap.pop()
    if heap:
        heap[0] = last
        i = 0
        while 2 * i + 1 < len(heap):
            child = 2 * i + 1
            if child + 1 < len(heap) and heap[child + 1] < heap[child]:
                child += 1
            if heap[i] <= heap[child]:
                break
            heap[i], heap[child] = heap[child], heap[i]
            i = child
    return smallest


heap = []
for value in [5, 3, 8, 1, 4]:
    push(heap, value)
print(heap)
print([pop(heap) for _ in range(5)])

[1, 3, 8, 5, 4]
[1, 3, 4, 5, 8]

The list isn't sorted, and it doesn't need to be: a heap only keeps the smallest in front, which is far less work than keeping everything in order.

heapq

Python's heapq module does the same in C, on a plain list:

>>> import heapq
>>> heap = [5, 3, 8, 1, 4]
>>> heapq.heapify(heap)
>>> heap
[1, 3, 8, 5, 4]
>>> heapq.heappush(heap, 2)
>>> heap[0]
1
>>> heapq.heappop(heap)
1
>>> heapq.heappop(heap)
2

heapify turns a list into a heap in place in O(n), cheaper than n pushes, which cost O(n log n). heappush and heappop are O(log n), and heap[0] is the smallest, O(1). Mind two things. Only heap[0] has a known place: heap[-1] isn't the largest, and heap[1] isn't always the second smallest. And heappop trusts that the list is a heap: on a list you never heapified, it takes whatever is at index 0.

heappush, heappop and heapify keep the smallest in front. For the largest first, push -x and negate what you pop, which works in every version of Python. Python 3.14 also has max-heap versions, such as heapq.heappush_max and heapq.heappop_max. Tuples compare item by item, so (priority, name) pairs come out by priority, with the name deciding ties.

The k largest, as they come

To keep the k largest values of a stream, keep a heap of at most k values. Its root is the smallest of the k largest so far: the k-th largest. When a value comes, push it, and if the heap now has k + 1 values, pop the smallest, which can't be among the k largest any more:

import heapq


def kth_largest(values, k):
    heap = []
    for value in values:
        heapq.heappush(heap, value)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]


scores = [40, 95, 12, 78, 66, 88, 51]
print(kth_largest(scores, 3), sorted(scores)[-3])

78 78

Each value costs O(log k), so a stream of n values costs O(n log k) time and O(k) memory, however long it is. And at every moment, heap[0] is the answer for the values seen so far, while a sorted list of all the values seen would cost O(n) for each new value, O(n²) in all. Once the heap is full, you can save more work: a value that isn't greater than heap[0] can't get in, so skip it, and for one that can, heapq.heapreplace(heap, value) pops the smallest and pushes the value in one sift.

The same idea merges k sorted lists: keep the next item of each list in a heap, as (value, list number, index), pop the smallest and push the next item of its list. That's O(N log k) for N items in all, and heapq.merge does it for you.