DSA course menu
Heaps and Priority Queues in Python (heapq Explained)
Learn the heap data structure: min-heaps and max-heaps, storing a complete tree in an array, sift up and sift down, building a heap in O(n), Python's heapq module, heap sort and top-k problems.
On this page
A priority queue always hands you the most important item next: the smallest number, the earliest deadline, the cheapest route. A heap is the data structure that makes this fast:
The heap property#
A min-heap is a complete binary tree where every parent is less than or equal to its children. So the smallest item is always at the root. (A max-heap flips this: every parent is greater or equal.)
1
/ \
3 2
/ \ /
7 4 5Siblings are not ordered relative to each other; only parent-child pairs are. That weaker rule is what makes heaps cheaper to maintain than a sorted list.
A tree stored in a list#
Because a heap is a complete tree, it fits in a plain list with no pointers. Store the nodes level by level; for the node at index i:
- left child at
2*i + 1 - right child at
2*i + 2 - parent at
(i - 1) // 2
The tree above is the list [1, 3, 2, 7, 4, 5].
Building a heap from scratch#
Insert: append to the end, then sift up: swap with the parent while smaller than it. Remove the minimum: take the root, move the last item to the root, then sift down: swap with the smaller child while bigger than it. Each moves along one root-to-leaf path: O(log n).
class MinHeap:
def __init__(self):
self.items = []
def push(self, value):
items = self.items
items.append(value)
i = len(items) - 1
while i > 0:
parent = (i - 1) // 2
if items[i] >= items[parent]:
break
items[i], items[parent] = items[parent], items[i]
i = parent
def pop(self):
items = self.items
smallest = items[0]
last = items.pop()
if items:
items[0] = last
i = 0
while True:
left, right = 2 * i + 1, 2 * i + 2
child = i
if left < len(items) and items[left] < items[child]:
child = left
if right < len(items) and items[right] < items[child]:
child = right
if child == i:
break
items[i], items[child] = items[child], items[i]
i = child
return smallest
h = MinHeap()
for x in [5, 3, 8, 1, 9, 2]:
h.push(x)
print("heap as a list:", h.items)
print([h.pop() for _ in range(6)])heap as a list: [1, 3, 2, 5, 9, 8]
[1, 2, 3, 5, 8, 9]Popping repeatedly returns the items in sorted order.
Python's heapq module#
heapq provides the same operations on an ordinary list (a min-heap):
import heapq
nums = [5, 3, 8, 1, 9, 2]
heapq.heapify(nums) # O(n), in place
print(nums)
heapq.heappush(nums, 0)
print(heapq.heappop(nums), heapq.heappop(nums))
print(nums[0]) # peek at the smallest without removing it[1, 3, 2, 5, 9, 8]
0 1
2heapify is O(n), not O(n log n): it sifts down from the last parent to the root, and most nodes sit near the bottom where sifting is short.
Priorities and max-heaps#
Store tuples to attach data to a priority; tuples compare by their first item. heapq only builds min-heaps, so negate the priority to get a max-heap:
import heapq
tasks = []
heapq.heappush(tasks, (2, "write report"))
heapq.heappush(tasks, (1, "fix production bug"))
heapq.heappush(tasks, (3, "reply to emails"))
while tasks:
priority, task = heapq.heappop(tasks)
print(priority, task)
scores = [70, 95, 82]
max_heap = [-s for s in scores]
heapq.heapify(max_heap)
print("highest score:", -max_heap[0])1 fix production bug
2 write report
3 reply to emails
highest score: 95Top k: the k largest items#
To keep the k largest items of a big stream, hold a min-heap of size k. Its root is the smallest of the current top k; any new item bigger than it replaces it. Total cost O(n log k), and only k items in memory:
import heapq
def top_k(nums, k):
heap = []
for x in nums:
if len(heap) < k:
heapq.heappush(heap, x)
elif x > heap[0]:
heapq.heapreplace(heap, x) # pop the root, push x
return sorted(heap, reverse=True)
data = [12, 3, 45, 7, 89, 23, 56, 1, 34]
print(top_k(data, 3))
print(heapq.nlargest(3, data), heapq.nsmallest(2, data))[89, 56, 45]
[89, 56, 45] [1, 3]The kth largest item is simply heap[0] at the end.
Heap sort#
Heapify, then pop everything: O(n log n), in place when written with a max-heap inside the array, and not stable.
import heapq
def heap_sort(items):
heap = items[:]
heapq.heapify(heap)
return [heapq.heappop(heap) for _ in range(len(heap))]
print(heap_sort([9, 4, 7, 1, 8, 2]))[1, 2, 4, 7, 8, 9]Heaps also power Dijkstra's shortest path, A*, merging k sorted lists and event scheduling.
Common mistakes#
Practice#
Practice exercise
Write merge_sorted_lists(lists) that merges k sorted lists into one sorted list using a heap that never holds more than k items.
Show hint
Push the first item of every list as (value, list_index, item_index). Pop the smallest, add it to the result, and push the next item from the same list if there is one.
Show solution
import heapq
def merge_sorted_lists(lists):
heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
heapq.heapify(heap)
result = []
while heap:
value, i, j = heapq.heappop(heap)
result.append(value)
if j + 1 < len(lists[i]):
heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
return result
print(merge_sorted_lists([[1, 4, 5], [1, 3, 4], [2, 6], []]))[1, 1, 2, 3, 4, 4, 5, 6]Each of the N items is pushed and popped once on a heap of at most k items: O(N log k). The list index in the tuple also breaks ties between equal values.
Key takeaway#
Next: a tree for words and prefixes, Tries.
Finished this lesson?