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?
Related lessons
Recommended DSA resources
- Cracking the Coding Interview by Gayle Laakmann McDowell
189 practice questions with solutions, for coding interview preparation.
Some links are affiliate links. If you buy through them, Laviyx Learn may earn a commission at no extra cost to you. As an Amazon Associate, we earn from qualifying purchases.