heapq Operations

Concepts covered: Heap & Top-K

A scheduler holding 50,000 waiting jobs keeps asking which job runs next. Calling min(jobs) reads all 50,000 every time, and re-sorting after each arrival reorders the whole list to answer a question about a single item. The heapq module keeps an ordinary list arranged so the smallest item always sits at index 0. After a push or a pop, a short walk through the list repairs that arrangement.

The heap rule: no parent is larger than its children

Picture index i as a parent whose 2 children live at indexes 2*i + 1 and 2*i + 2. A list is a heap when every parent is less than or equal to both children. That rule puts the minimum at index 0 and says nothing about the order of the rest. Watch the list after each heappush:

import heapq

heap = []
for minutes in [7, 2, 9, 4, 1]:
    heapq.heappush(heap, minutes)
    print(heap)

print("smallest:", heap[0])
print("popped:", heapq.heappop(heap))
print("left:", heap)

Output:

[7]
[2, 7]
[2, 7, 9]
[2, 4, 9, 7]
[1, 2, 9, 7, 4]
smallest: 1
popped: 1
left: [2, 4, 9, 7]

The final list [1, 2, 9, 7, 4] is out of order, since 9 sits before 7 and 4, yet heap[0] holds the minimum. When 1 arrived it was appended at index 4, compared with its parent at index 1, swapped, then compared with index 0 and swapped again. A pop runs the other direction: the last item moves to index 0 and sinks, swapping with its smaller child at each level.

Those walks stay short because each level of the tree holds twice as many items as the level above it: 1, then 2, then 4, so 20 levels hold over 1,000,000 items. A push or pop on a 1,000,000-item heap visits roughly 20 levels at most, where min() reads every item. For data already sitting in a list, heapq.heapify(items) rearranges it into a heap in place, in a single pass.

Tuples carry the priority

Heaps compare items with <, so a tuple (priority, job) orders by priority first and compares job names only on a tie. Lower numbers leave first. The 2 jobs below share priority 1; decide which pops first, then play it.

import heapq

task_queue = []
heapq.heappush(task_queue, (2, "load_orders"))
heapq.heappush(task_queue, (1, "fix_schema"))
heapq.heappush(task_queue, (3, "rebuild_docs"))
heapq.heappush(task_queue, (1, "backfill_events"))

print("Processing tasks by priority:")
while task_queue:
    priority, task = heapq.heappop(task_queue)
    print(f"  [{priority}] {task}")

Processing tasks by priority:
  [1] backfill_events
  [1] fix_schema
  [2] load_orders
  [3] rebuild_docs

Both priority-1 jobs tied, so Python compared the names, and 'backfill_events' sorts before 'fix_schema'. That fallback fails when the second item can't be compared, such as a dict describing the job. Put a sequence number between the priority and the payload: no 2 entries share a number, so Python never reaches the payload, and equal priorities leave in arrival order.

import heapq

jobs = []
heapq.heappush(jobs, (1, {"name": "load_orders"}))
try:
    heapq.heappush(jobs, (1, {"name": "load_users"}))
except TypeError as err:
    print("TypeError:", err)

jobs = []
for seq, name in enumerate(["load_orders", "load_users"]):
    heapq.heappush(jobs, (1, seq, {"name": name}))
print(heapq.heappop(jobs))

Output:

TypeError: '<' not supported between instances of 'dict' and 'dict'
(1, 0, {'name': 'load_orders'})

Largest first, and the top k

heapq builds min-heaps only. For largest-first order, push each number negated and negate it again on the way out, because the most negative stored value is the largest original.

import heapq

max_heap = []
for v in [5, 3, 8, 1, 9]:
    heapq.heappush(max_heap, -v)

print("Pop in descending order:")
while max_heap:
    print(-heapq.heappop(max_heap))

Pop in descending order:
9
8
5
3
1

When you need a top k once and have no queue to keep alive, heapq.nlargest(k, items) and heapq.nsmallest(k, items) answer in a single call. Both take key= the way sorted() does. The last line gets the same answer by sorting all 5 values to keep 2.

import heapq

partition_mb = {"2024-01": 120, "2024-02": 950, "2024-03": 40,
                "2024-04": 610, "2024-05": 300}

def size_of(name):
    return partition_mb[name]

print(heapq.nlargest(2, partition_mb, key=size_of))
print(heapq.nsmallest(2, partition_mb.values()))
print(sorted(partition_mb.values())[:2])

Output:

['2024-02', '2024-04']
[40, 120]
[40, 120]

Only index 0 has a fixed meaning. The largest item always sits among the leaves, the slots with no children. In a list of n items those run from index n // 2 to n - 1, and nothing says which leaf holds the maximum. In [1, 2, 9, 7, 4] the maximum, 9, is at index 2. Meanwhile heap[-1] holds 4, the item that 1 displaced on its climb to the top. Asking a min-heap for its largest item still means scanning half the list.

What does this program print?

import heapq

h = []
for n in [6, 4, 9, 2]:
    heapq.heappush(h, n)
print(h[0], h[-1])

• 2 9, because index 0 holds the minimum and the last index holds the maximum
• 6 2, because a heap list keeps items in the order they were pushed
• 2 6, because 2 climbed to index 0 and pushed 6 down into the last slot
• 2 4, because a heap list stays fully sorted from smallest to largest

About this lesson

This section is part of the Collections: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.