HEAP-4

heapq and priority queues

HeapsLesson 4 of 615 min

Python’s heapq module is the heap from lessons 1 to 3, written in C, working on an ordinary list. This lesson shows how to use it as a priority queue: a queue where the most urgent item leaves first, whatever order the items arrived in.

The heapq module

heapq has no class. Its functions take a plain list and keep it in heap order, smallest at a[0]. Keep using the module’s functions on that list; a plain append or sort breaks the heap.

The main heapq functions
CallWhat it doesCost
heapify(a)turn the list into a heap, in placeO(n)
heappush(a, x)add xO(log n)
heappop(a)remove and return the smallestO(log n)
a[0]look at the smallestO(1)
heapreplace(a, x)pop, then push x (one sift)O(log n)
heappushpop(a, x)push x, then pop (often no sift at all)O(log n)
nsmallest(k, xs), nlargest(k, xs)the k smallest or largest, in orderO(n log k)
merge(*lists)merge sorted inputs, lazilyO(n log k)
import heapq

a = [5, 3, 8, 1]
heapq.heapify(a)  # a is now [1, 3, 8, 5]
heapq.heappush(a, 2)
heapq.heappop(a)  # 1
heapq.heappop(a)  # 2

Max-heaps

heapq builds min-heaps. For numbers, the usual way to get a max-heap is to store each value negated: the smallest of the negated values is the largest of the real ones. Push -x, and negate again when you pop.

h = []
for x in [5, 3, 8]:
    heapq.heappush(h, -x)
-heapq.heappop(h)  # 8, the largest

Python 3.14 (the version on the practice pages) also added max-heap functions: heapify_max, heappush_max, heappop_max, heapreplace_max and heappushpop_max. Older Pythons only have negation.

Priority queues and ties

To queue tasks by priority, push tuples (priority, task). Tuples compare item by item, so the smallest priority comes out first. Two problems appear when priorities tie:

  • Python then compares the tasks themselves. If they are objects that cannot be compared (say, two dicts), you get a TypeError.
  • Even if they can be compared, the result is not “first come, first served”. A heap is not stable: items with equal priority come out in no particular order.

The fix for both is a counter that goes up by one with every push: store (priority, count, task). Counts are all different, so the task is never compared, and among equal priorities the earliest arrival wins. Step through the figure.

FIG HEAP-4.1A priority queue with ties and removal
1/10

Each entry is a tuple (priority, count, task). Add write with priority 3. The count is 0: it numbers the entries in the order they arrive.

Step 1 of 10

Changing and removing

A heap cannot find an item quickly, so it cannot remove one from the middle quickly either. The usual answer, given in the heapq documentation, is lazy removal: keep a dict from each task to its entry, and to remove a task, mark its entry as removed and leave it in the heap. When a marked entry later reaches the top, pop throws it away and tries again. To change a task’s priority, remove it and push a new entry.

import heapq, itertools

heap, entry_of = [], {}
counter = itertools.count()
REMOVED = None

def add(task, priority):
    if task in entry_of:
        remove(task)
    e = [priority, next(counter), task]
    entry_of[task] = e
    heapq.heappush(heap, e)

def remove(task):
    e = entry_of.pop(task)
    e[-1] = REMOVED  # mark it; leave it in the heap

def pop():
    while heap:
        priority, count, task = heapq.heappop(heap)
        if task is not REMOVED:
            del entry_of[task]
            return task
    raise KeyError('pop from an empty priority queue')

Entries are lists rather than tuples so that the mark can be written into them. Each removed entry is popped once, later, so every operation is still O(log n) on average; the heap can hold up to twice as many entries as there are live tasks before they are cleared out. Practice HEAP-P3 asks you to build this queue (with heapq switched off, so with your own heap).

Other tools

  • queue.PriorityQueue wraps heapq with a lock, for passing work between threads. In a single thread, use heapq directly: it is faster.
  • For items that cannot be compared, the dataclasses module can make a class that compares only its priority field: @dataclass(order=True) with item: Any = field(compare=False).

HEAP-4 counts as completed once you step through the figure to the last step.