Heaps and priority queues explained: top-k, merging and task scheduling

A heap is the data structure behind priority queues: you can always take out the smallest (or largest) item quickly, without keeping the whole collection sorted. It shows up in interviews ("top k", "merge k lists", "running median") and in real systems such as schedulers and search. The code, tests and timings below were run with Python 3 before publishing.

How a min-heap works

A min-heap is a binary tree where every parent is smaller than or equal to its children. So the smallest item is always at the top. It is stored in a plain list: the children of position i are at 2i+1 and 2i+2, and the parent is at (i-1)//2.

Operation	Cost	How
Peek at the smallest	O(1)	Read position 0
Push	O(log n)	Add at the end, then swap upwards while smaller than the parent
Pop the smallest	O(log n)	Swap the last item to the top, remove the old top, swap downwards
Build from n items	O(n)	Fix the lower half from the bottom up ("heapify")

A heap is not fully sorted. Only the parent-child rule holds, which is exactly why it is cheaper than sorting when you only need the next smallest item.

A working example

Save as heaps.py:

import heapq
import random
import time


class MinHeap:
    """A binary min-heap stored in a list: children of i are 2i+1 and 2i+2."""

    def __init__(self, items=()):
        self.a = list(items)
        for i in range(len(self.a) // 2 - 1, -1, -1):  # heapify in O(n)
            self._down(i)

    def __len__(self):
        return len(self.a)

    def peek(self):
        return self.a[0]

    def push(self, x):
        self.a.append(x)
        self._up(len(self.a) - 1)

    def pop(self):
        a = self.a
        a[0], a[-1] = a[-1], a[0]
        smallest = a.pop()
        if a:
            self._down(0)
        return smallest

    def _up(self, i):
        a = self.a
        while i > 0:
            parent = (i - 1) // 2
            if a[parent] <= a[i]:
                break
            a[parent], a[i] = a[i], a[parent]
            i = parent

    def _down(self, i):
        a, n = self.a, len(self.a)
        while True:
            left, right, small = 2 * i + 1, 2 * i + 2, i
            if left < n and a[left] < a[small]:
                small = left
            if right < n and a[right] < a[small]:
                small = right
            if small == i:
                return
            a[i], a[small] = a[small], a[i]
            i = small


def top_k_largest(items, k):
    """Keep a min-heap of size k: O(n log k) time, O(k) memory."""
    heap = []
    for x in items:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:
            heapq.heapreplace(heap, x)
    return sorted(heap, reverse=True)


def merge_sorted(lists):
    """Merge k sorted lists: the heap always holds one front item per list."""
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)
    out = []
    while heap:
        value, i, j = heapq.heappop(heap)
        out.append(value)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return out


class TaskQueue:
    """Lowest priority number first; a counter keeps equal priorities in insertion order."""

    def __init__(self):
        self._heap, self._n = [], 0

    def add(self, priority, task):
        heapq.heappush(self._heap, (priority, self._n, task))
        self._n += 1

    def next(self):
        return heapq.heappop(self._heap)[2]

    def __len__(self):
        return len(self._heap)


if __name__ == "__main__":
    h = MinHeap([9, 4, 7, 1, 8, 2])
    h.push(0)
    print("MinHeap pops :", [h.pop() for _ in range(len(h))])

    q = TaskQueue()
    for p, t in [(2, "send email"), (1, "fix outage"), (2, "write report"), (3, "tidy desk")]:
        q.add(p, t)
    print("task order   :", [q.next() for _ in range(len(q))])

    print("merge 3 lists:", merge_sorted([[1, 4, 9], [2, 3, 10], [5, 6, 7, 8]]))

    rnd = random.Random(1)
    data = [rnd.randint(0, 10**9) for _ in range(300_000)]
    t = time.perf_counter(); a = top_k_largest(data, 10); t1 = (time.perf_counter() - t) * 1000
    t = time.perf_counter(); b = sorted(data, reverse=True)[:10]; t2 = (time.perf_counter() - t) * 1000
    t = time.perf_counter(); c = heapq.nlargest(10, data); t3 = (time.perf_counter() - t) * 1000
    print(f"top 10 of 300,000: same answer = {a == b == c}")
    print(f"  heap of size k (ours)   {t1:7.1f} ms")
    print(f"  heapq.nlargest          {t3:7.1f} ms")
    print(f"  sorted(...)[:10]        {t2:7.1f} ms")

Output of python3 heaps.py (exact times depend on the machine):

MinHeap pops : [0, 1, 2, 4, 7, 8, 9]
task order   : ['fix outage', 'send email', 'write report', 'tidy desk']
merge 3 lists: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
top 10 of 300,000: same answer = True
  heap of size k (ours)      38.0 ms
  heapq.nlargest             12.7 ms
  sorted(...)[:10]          192.5 ms

What the code shows

• MinHeap is the structure written by hand. Popping all items gives them in sorted order, which is also the idea behind heap sort.
• top_k_largest keeps a min-heap of size k. For each new item bigger than the smallest of the k kept so far, it replaces that one. Time O(n log k), memory O(k), so it works on a stream you cannot hold in memory.
• merge_sorted merges k sorted lists by keeping one front item per list in the heap. Each step costs O(log k), so merging n items total costs O(n log k).
• TaskQueue stores (priority, counter, task). The counter breaks ties, so equal priorities come out in the order they were added, and Python never has to compare two tasks directly.

Reading the timings

For the 10 largest of 300,000 numbers:

• sorted(data, reverse=True)[:10] took 192 ms, because it sorts everything to keep 10 items.
• Our heap of size 10 took 38 ms, about 5 times faster.
• heapq.nlargest took 12.7 ms, faster again because it is written in C. Our hand-written version is plain Python.

If k is close to n, sorting is just as good. The heap wins when k is small compared to n.

Python's heapq

Python ships heapq, a min-heap on a list: heappush, heappop, heapify, heapreplace, nlargest, nsmallest. Two things to remember:

• It is a min-heap only. For a max-heap, store negated numbers (or negated priorities).
• Items must be comparable. For tuples, Python compares the first element, then the next, which is why the counter above matters.

Tests

Save as test_heaps.py:

import heapq
import random
import unittest

from heaps import MinHeap, TaskQueue, merge_sorted, top_k_largest


class HeapTests(unittest.TestCase):
    def test_minheap_pops_in_order(self):
        rnd = random.Random(3)
        for _ in range(200):
            data = [rnd.randint(-20, 20) for _ in range(rnd.randint(0, 50))]
            h = MinHeap(data)
            self.assertEqual([h.pop() for _ in range(len(h))], sorted(data))

    def test_push_then_pop(self):
        h = MinHeap()
        for x in [5, 3, 8, 1]:
            h.push(x)
        self.assertEqual(h.peek(), 1)
        self.assertEqual([h.pop() for _ in range(4)], [1, 3, 5, 8])

    def test_top_k_matches_sorted(self):
        rnd = random.Random(5)
        for _ in range(100):
            data = [rnd.randint(0, 100) for _ in range(rnd.randint(0, 80))]
            for k in (1, 3, 10):
                self.assertEqual(top_k_largest(data, k), sorted(data, reverse=True)[:k])

    def test_merge_sorted(self):
        rnd = random.Random(9)
        lists = [sorted(rnd.randint(0, 50) for _ in range(rnd.randint(0, 10))) for _ in range(6)]
        self.assertEqual(merge_sorted(lists), sorted(x for l in lists for x in l))
        self.assertEqual(merge_sorted([[], []]), [])

    def test_task_queue_is_stable_for_equal_priority(self):
        q = TaskQueue()
        for t in ["a", "b", "c"]:
            q.add(1, t)
        self.assertEqual([q.next() for _ in range(3)], ["a", "b", "c"])


if __name__ == "__main__":
    unittest.main()

$ python3 -m unittest -v
Ran 5 tests in 0.025s

OK

Each structure is checked against the simple way of getting the same answer: sorted() for the heap, top-k and merge. Random inputs include empty lists and many duplicates.

When to use a heap

• Top k or k smallest of a large or streaming collection.
• Merge k sorted sources (log files, database shards, search results).
• Scheduling by priority or by earliest time (timers, job queues, Dijkstra's shortest path).
• Running median: one max-heap for the lower half and one min-heap for the upper half.

If you need to look items up by value, or remove arbitrary items quickly, a heap is the wrong tool: finding an item is O(n).

Common mistakes

• Thinking heap[1] is the second smallest. Only heap[0] is guaranteed.
• Sorting everything to get the top few items.
• Expecting heapq to be a max-heap.
• Pushing tuples whose second element cannot be compared, with no tie-breaker.
• Changing an item's priority inside the heap without fixing the heap. Push a new entry and skip the old one when you pop it.

Main idea: a heap gives O(log n) push and pop and O(1) access to the smallest item. Use it when you repeatedly need the next best item, especially top-k and merging, and use heapq in real code.

---

Interview prep and DSA series, Part 3 (more algorithms): 2 of 5. Previous: Sorting algorithms explained: insertion, merge and quick sort with timings.