Python heapq

How heapq stores a heap in a list

The heapq module works on a plain Python list. It does not make a new class. The module keeps the smallest item at index 0.

The list order is a heap order, not a sorted order. The parent of index i is at (i - 1) // 2. The children of index i are at 2 * i + 1 and 2 * i + 2.

You can call the functions on any list. The list must be a valid heap before you call heappop or heappush.

The four core functions

heapify turns a list into a heap in place. heappush adds one item and keeps the heap order. heappop removes and returns the smallest item. heapreplace removes the smallest item and adds a new item in one step.

The costs are simple. heappush and heappop are O(log n). heapify is O(n). Reading the item at index 0 is O(1).

The code below builds a heap and changes it. The output shows the smallest item first. heapreplace returns the old smallest item, which is 1.

import heapq

nums = [5, 1, 3, 2, 4]
heapq.heapify(nums)
print(nums[0])
heapq.heappush(nums, 0)
print(heapq.heappop(nums))
print(heapq.heapreplace(nums, 6))
print(nums)

A task queue with tuples

A common use is a task queue. Store a tuple of (priority, name) in the heap. Python compares the first item first, so the lowest priority number comes out first.

If two tasks have the same priority, Python compares the names. The names are strings, so the compare works. The code below adds three tasks and pops them in priority order.

The output is ('a', 'write') first, then ('b', 'test'), then ('c', 'ship').

import heapq

tasks = []
heapq.heappush(tasks, (2, "test"))
heapq.heappush(tasks, (1, "write"))
heapq.heappush(tasks, (3, "ship"))
while tasks:
    print(heapq.heappop(tasks))

How to build a max heap

heapq has no max heap function. The usual trick is to negate the key. Push -value and the smallest item is the largest value.

When you pop, negate the result again. The code below keeps the three largest values on top. The output is 9, then 7, then 4.

import heapq

values = [4, 9, 7]
heap = []
for v in values:
    heapq.heappush(heap, -v)
while heap:
    print(-heapq.heappop(heap))

nlargest, nsmallest, and merge

The module has helper functions for common jobs. nlargest(k, items) returns the k largest items. nsmallest(k, items) returns the k smallest items. Both return a sorted list.

merge(*iterables) joins sorted inputs into one sorted stream. It does not load all items at once. This makes it useful for large files.

The code below shows all three helpers. The output is [9, 7] for the largest, [1, 2] for the smallest, and 1 2 3 4 for the merge.

import heapq

print(heapq.nlargest(2, [4, 9, 1, 7, 2]))
print(heapq.nsmallest(2, [4, 9, 1, 7, 2]))
for x in heapq.merge([1, 3], [2, 4]):
    print(x)

Common mistakes and fixes

Do not append to a heap with list.append. This breaks the heap order. Fix: use heapq.heappush.

Do not change a tuple after the push. The heap order becomes wrong. Fix: push a new tuple and ignore the old one.

Do not push items that cannot compare. A tuple with a dict as the second item fails when priorities tie. Fix: add a counter as the second item.

Do not use sorted() for a streaming queue. sorted() costs O(n log n) on each new item. Fix: use heappush and heappop, which cost O(log n).

Common questions

Is heapq a queue or a list?

heapq is a module that works on a plain list. The list is the storage, and the module keeps the heap order.

How do I make a max heap?

Negate the key before the push. Negate the value again after the pop.

When is heapq faster than sorting?

heapq is faster when you add and remove items over time. Sorting is faster when you need one full order of all items.