Min-heap (priority queue) in Python

How it works

1. The list is read as a tree: the children of index i are at 2 * i + 1 and 2 * i + 2, and its parent is at (i - 1) // 2.
2. push adds the new value at the end, then swaps it with its parent while the parent is bigger: it sifts up.
3. pop takes heap[0], the smallest value, moves the last value to the top, then swaps it with its smaller child until both children are bigger: it sifts down.
4. Python's heapq module does the same with heapq.heappush and heapq.heappop, on a plain list.

What to watch

• Watch 1 sift up from the bottom to the top when it is pushed.
• Watch the tree after each push: every parent is at most its children, but the list itself is not sorted.
• Watch child pick the smaller of the two children before each swap down.

The code

def push(heap, value):
    heap.append(value)                  # add at the bottom
    i = len(heap) - 1
    while i > 0 and heap[(i - 1) // 2] > heap[i]:
        parent = (i - 1) // 2
        heap[i], heap[parent] = heap[parent], heap[i]
        i = parent                      # move up the tree

def pop(heap):
    smallest, last = heap[0], heap.pop()
    if heap:
        heap[0] = last                  # last value to the top
        i = 0
        while 2 * i + 1 < len(heap):
            child = 2 * i + 1
            if child + 1 < len(heap) and heap[child + 1] < heap[child]:
                child += 1              # the smaller child
            if heap[i] <= heap[child]:
                break
            heap[i], heap[child] = heap[child], heap[i]
            i = child                   # move down the tree
    return smallest

heap = []
for value in [7, 3, 9, 1, 5, 2]:
    push(heap, value)
print("Heap:", heap)
print("Taken out:", pop(heap), pop(heap), pop(heap))

Heap: [1, 3, 2, 7, 5, 9]
Taken out: 1 2 3

Open it in the editor · Python Program to Implement a Priority Queue, explained