Trees & Heaps -- Inventing Fast Lookup

You keep a collection of numbers. Items are inserted, deleted, and searched constantly. What is the best data structure?

Structure insert find delete
plain list O(1) append O(n) scan O(n) find + O(n) shift
sorted list O(n) shift O(log n) binary search O(n) shift
dictionary O(1) O(1) O(1) -- but heavy memory, no ordering

Each structure is fast at some operations and slow at others. In this chapter you will invent a structure that does all three in O(log n) while keeping the data ordered -- the binary search tree -- and then a specialized cousin that powers every task scheduler: the heap.

Part 1: Warm-Up -- Trees That Compute

A tree is nodes connected downward: each node holds a value and has up to two children, left and right.

Here is the arithmetic expression 2 * 3 + 4 drawn as a tree:

        +
       / \
      *   4
     / \
    2   3

Leaves are numbers; internal nodes are operators. This is exactly how compilers see your code (Python, SQL, C++ -- everything is parsed into trees).

The Node class and a tree printer are provided -- run this cell as-is.

Exercise 1.1 — Evaluate an Expression Tree

Write evaluate(node) that computes the value of an expression tree.

Before you code:

Support +, -, *, /.

evaluate(expr)   # 10   (2*3 + 4)
Provided — Node class and tree printer
class Node:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def print_tree(node, indent=0):
    """Prints the tree sideways: root at the left, right child on top."""
    if node is None:
        return
    print_tree(node.right, indent + 1)
    print("     " * indent + str(node.val))
    print_tree(node.left, indent + 1)


# The tree for 2 * 3 + 4:
expr = Node("+", Node("*", Node(2), Node(3)), Node(4))
print_tree(expr)
Provided — test cases
assert evaluate(expr) == 10
assert evaluate(Node(42)) == 42
assert evaluate(Node("-", Node(10), Node(4))) == 6
print("evaluate: OK")
Hint 1

Check whether node.val is numeric (a leaf). If so, return it directly. Otherwise, recursively evaluate the left and right children first, then apply the operator stored in node.val.

Hint 2

Use isinstance(node.val, (int, float)) to detect leaves. For operators, compute l = evaluate(node.left) and r = evaluate(node.right), then branch on node.val being "+", "-", "*", or "/".

Solution
def evaluate(node):
    if isinstance(node.val, (int, float)):
        return node.val
    l = evaluate(node.left)
    r = evaluate(node.right)
    if node.val == "+":
        return l + r
    elif node.val == "-":
        return l - r
    elif node.val == "*":
        return l * r
    elif node.val == "/":
        return l / r

Exercise 1.2 — Build a Tree Yourself

Build the tree for (2 + 3) * (10 - 4) using Node, print it with print_tree, and check that evaluate returns 30.

Notice: the tree needs no parentheses. The structure is the precedence.

Provided — test
print_tree(expr2)
assert evaluate(expr2) == 30
print("expr2: OK")
Hint 1

The root is "*". Its left child is "+" with children 2 and 3. Its right child is "-" with children 10 and 4.

Hint 2

expr2 = Node("*", Node("+", Node(2), Node(3)), Node("-", Node(10), Node(4)))

Solution
expr2 = Node("*", Node("+", Node(2), Node(3)), Node("-", Node(10), Node(4)))

Part 2: The Binary Search Tree

Back to our fast-lookup problem. Here is the one rule that makes a tree searchable:

BST rule: everything in a node's left subtree is smaller than the node's value; everything in the right subtree is larger.

Play with the animation before coding: BST Visualization. Insert 10, 12, 1, 11, 13, 3 and watch where each lands.

Exercise 2.1 — Insert

Write insert(node, val) that inserts val into the tree rooted at node and returns the root of the resulting tree.

Think recursively:

root = None
for v in [10, 12, 1, 11, 13, 3]:
    root = insert(root, v)
print_tree(root)

Compare the printed tree with what the animation showed you.

Provided — tests
assert root.val == 10
assert root.left.val == 1
assert root.right.val == 12
assert root.left.right.val == 3
assert root.right.left.val == 11
assert root.right.right.val == 13
print("insert: OK")
Hint 1

Base case: if node is None, return Node(val). Recursive case: if val < node.val, recurse left; otherwise recurse right. Always return node at the end so the parent can reattach.

Hint 2
def insert(node, val):
    if node is None:
        return Node(val)
    if val < node.val:
        node.left = insert(node.left, val)
    else:
        node.right = insert(node.right, val)
    return node
Solution
def insert(node, val):
    if node is None:
        return Node(val)
    if val < node.val:
        node.left = insert(node.left, val)
    else:
        node.right = insert(node.right, val)
    return node

Exercise 2.2 — Find

Write contains(node, val) that returns True if val is in the tree.

The payoff: at each node you compare once and discard half the tree -- the same halving you invented in the Binary Search chapter. That is why find is O(log n)... if the tree is balanced (more on that soon).

contains(root, 11)   # True
contains(root, 7)    # False
Provided — tests
assert contains(root, 11) == True
assert contains(root, 10) == True
assert contains(root, 13) == True
assert contains(root, 7)  == False
assert contains(None, 5)  == False
print("contains: OK")
Hint 1

Base case: if node is None, return False. If val == node.val, return True. Otherwise go left or right depending on whether val is smaller or larger.

Solution
def contains(node, val):
    if node is None:
        return False
    if val == node.val:
        return True
    if val < node.val:
        return contains(node.left, val)
    else:
        return contains(node.right, val)

Exercise 2.3 — Walk the Tree in Order

Write in_order(node) that returns a list of all values using this recipe: first everything in the left subtree, then the node itself, then the right subtree.

in_order(root)   # ?

Run it. Look carefully at the output. What did you just get for free? (You have invented tree sort: insert n values, walk once, sorted output.)

Provided — tests
result = in_order(root)
print(result)
assert result == [1, 3, 10, 11, 12, 13]

# tree sort on random data:
values = random.sample(range(1000), 50)
r = None
for v in values:
    r = insert(r, v)
assert in_order(r) == sorted(values)
print("in_order: OK -- the BST keeps data sorted at all times")
Hint 1

Base case: if node is None, return an empty list []. Recursive case: return in_order(node.left) + [node.val] + in_order(node.right).

Solution
def in_order(node):
    if node is None:
        return []
    return in_order(node.left) + [node.val] + in_order(node.right)

Exercise 2.4 — Find the Minimum

Where does the smallest value in a BST live? Trace it on the picture.

Write find_min(node) that returns the smallest value in a non-empty tree. No comparisons needed -- just walking.

find_min(root)   # 1
Provided — tests
assert find_min(root) == 1
assert find_min(root.right) == 11
assert find_min(Node(42)) == 42
print("find_min: OK")
Hint 1

The minimum is the leftmost node. Keep going left until there is no left child: while node.left: node = node.left. Return node.val.

Solution
def find_min(node):
    while node.left:
        node = node.left
    return node.val

Exercise 2.5 — ...and the Maximum

Mirror image: write find_max(node) returning the largest value. Same rule -- no comparisons, just walk. Both of these run in O(height): on a balanced tree of a million values, about 20 steps to either extreme.

Provided — tests
assert find_max(root) == 13
assert find_max(root.left) == 3
assert find_max(Node(42)) == 42
print("find_max: OK")
Hint 1

Mirror of find_min: keep going right until there is no right child.

Solution
def find_max(node):
    while node.right:
        node = node.right
    return node.val

Exercise 2.6 — Delete

Deletion is the tricky one.

Write delete(node, val) that removes val and returns the root of the resulting subtree. Think about what happens in each case:

root = delete(root, 1)     # leaf
root = delete(root, 12)    # two children
in_order(root)             # [3, 10, 11, 13]
Provided — tests
# rebuild a fresh tree
root = None
for v in [10, 12, 1, 11, 13, 3]:
    root = insert(root, v)

root = delete(root, 1)     # case: node with one child (1 has right child 3)
assert in_order(root) == [3, 10, 11, 12, 13]

root = delete(root, 13)    # case: leaf
assert in_order(root) == [3, 10, 11, 12]

root = delete(root, 10)    # case: two children (the root itself!)
assert in_order(root) == [3, 11, 12]

root = delete(root, 99)    # deleting a missing value changes nothing
assert in_order(root) == [3, 11, 12]
print("delete: OK")

# stress test: delete everything in random order
values = random.sample(range(500), 100)
r = None
for v in values:
    r = insert(r, v)
random.shuffle(values)
remaining = sorted(values)
for v in values:
    r = delete(r, v)
    remaining.remove(v)
    assert in_order(r) == remaining, f"tree broken after deleting {v}"
print("delete stress test: OK")
Hint 1

First, search for the node as in contains. When you find it, handle the three cases. For the two-children case: node.val = find_min(node.right), then node.right = delete(node.right, node.val).

Hint 2

For the leaf case, return None. For one child, return whichever child exists (left or right). The successor is guaranteed to have at most one child because it is the leftmost node in the right subtree -- it cannot have a left child.

Hint 3
def delete(node, val):
    if node is None:
        return None
    if val < node.val:
        node.left = delete(node.left, val)
    elif val > node.val:
        node.right = delete(node.right, val)
    else:
        # found it
        if node.left is None:
            return node.right
        if node.right is None:
            return node.left
        # two children
        node.val = find_min(node.right)
        node.right = delete(node.right, node.val)
    return node
Solution
def delete(node, val):
    if node is None:
        return None
    if val < node.val:
        node.left = delete(node.left, val)
    elif val > node.val:
        node.right = delete(node.right, val)
    else:
        # Found the node to delete
        if node.left is None:
            return node.right
        if node.right is None:
            return node.left
        # Two children: replace with in-order successor
        node.val = find_min(node.right)
        node.right = delete(node.right, node.val)
    return node

Exercise 2.7 — The Catch -- Balance

Insert 1, 2, 3, 4, 5, 6, 7 in that order and print the tree. What shape do you get?

Every insert went right: the "tree" is a chain -- a glorified linked list, and find is back to O(n). Sorted input is the worst input for a plain BST.

Self-balancing trees (AVL, Red-Black) fix this by rotating nodes on insert to keep the height at log n. Watch one rebalance itself: Red-Black Tree Visualization -- insert 1...10 in order and watch it stay shallow. (Databases use B-Trees -- the same idea, made disk-friendly.)

Provided — demonstration
chain = None
for v in [1, 2, 3, 4, 5, 6, 7]:
    chain = insert(chain, v)
print_tree(chain)
Solution

Quick Check 2.8 — BST Search Complexity

You insert the numbers 1, 2, 3, 4, 5 (in that order) into an empty BST. What does the tree look like, and how many comparisons does finding the value 5 require?

Hint

When you insert 2 into a tree containing only 1, where does 2 go? Then where does 3 go relative to 2?

Reasoning

Each new value is larger than all existing values, so it always goes to the right child. The tree becomes 1->2->3->4->5, a chain with no left children -- effectively a linked list. Finding 5 requires walking all 5 nodes. This is the worst case for BSTs: O(n) instead of the hoped-for O(log n). Balanced BST variants (AVL, Red-Black) prevent this by rebalancing after insertions.

Part 3: The Scheduler Problem & The Heap

You are building a task scheduler (like Unix cron). Tasks arrive at any time:

The scheduler repeatedly needs just one thing: the earliest task. So the operations are: add(task), pop_earliest(). Nothing else. No search, no arbitrary delete.

Before reading on, rate each strategy (fill the ?s in your head):

Strategy add pop earliest
unsorted list O(1) O(?) scan
keep list sorted O(?) O(1)
BST O(log n) O(log n)
one slot per second for 1000 days O(1) O(?) -- and how much memory?

The BST wins -- but it maintains total order, and we only ever ask for the min. Can a structure that promises less do the job with less work? It exists: the min-heap.

The Heap Shape

A min-heap is a binary tree with one weak promise:

Heap rule: every node is <= its children. (Nothing about left vs right!)

So the minimum is always... where?

Because the shape is always a "filled top-down, left-to-right" tree, we don't need Node objects at all -- we can store it in a plain list:

        1
      /   \
     3     2          stored as:  [1, 3, 2, 8, 5]
    / \
   8   5

Exercise 3.1 — The Index Arithmetic

For the node at list index i, figure out (try it on the picture above):

Write the three functions.

parent(3)       # 1   (8's parent is 3)
left_child(1)   # 3
right_child(0)  # 2
Provided — tests
assert parent(3) == 1 and parent(4) == 1 and parent(1) == 0 and parent(2) == 0
assert left_child(0) == 1 and right_child(0) == 2
assert left_child(1) == 3 and right_child(1) == 4
print("index arithmetic: OK")
Hint 1

Try the formulas on the picture. For 0-indexed arrays: parent(i) = (i - 1) // 2, left_child(i) = 2*i + 1, right_child(i) = 2*i + 2.

Solution
def parent(i):
    return (i - 1) // 2

def left_child(i):
    return 2 * i + 1

def right_child(i):
    return 2 * i + 2

Exercise 3.2 — Push

To add a value, append it at the end of the list (this keeps the shape correct). But the heap rule might now be broken -- the new value could be smaller than its parent. How would you fix this?

Write heap_push(heap, val) (in place).

h = [1, 3, 2, 8, 5]
heap_push(h, 0)
h   # [0, 3, 1, 8, 5, 2]  -- 0 climbed to the top

How many swaps can a push cost at most, for a heap of n items? (The tree has $\log_2(n)$ levels -- that is the whole point.)

Provided — tests
h = [1, 3, 2, 8, 5]
heap_push(h, 0)
assert h == [0, 3, 1, 8, 5, 2]

h = []
for v in [5, 3, 8, 1]:
    heap_push(h, v)
assert h[0] == 1, "minimum must be at the root"
print("heap_push: OK")
Hint 1

Append the value, then set i = len(heap) - 1. While i > 0 and heap[i] < heap[parent(i)], swap them and set i = parent(i).

Solution
def heap_push(heap, val):
    heap.append(val)
    i = len(heap) - 1
    while i > 0 and heap[i] < heap[parent(i)]:
        heap[i], heap[parent(i)] = heap[parent(i)], heap[i]
        i = parent(i)

Exercise 3.3 — Pop the Minimum

The minimum is heap[0]. Removing it leaves a hole at the root. To keep the shape, move the last element into the root position. But now the heap rule might be broken -- how would you fix it?

Write heap_pop(heap) that removes and returns the minimum. In place.

Edge case to think about: why must you swap with the smaller child? What if a node has only a left child?

h = [1, 3, 2, 8, 5]
heap_pop(h)   # 1
heap_pop(h)   # 2
heap_pop(h)   # 3
Provided — tests
h = [1, 3, 2, 8, 5]
assert heap_pop(h) == 1
assert heap_pop(h) == 2
assert heap_pop(h) == 3
assert heap_pop(h) == 5
assert heap_pop(h) == 8
assert h == []

# stress: pushes and pops always yield values in sorted order
values = [random.randint(0, 1000) for _ in range(500)]
h = []
for v in values:
    heap_push(h, v)
popped = [heap_pop(h) for _ in range(len(values))]
assert popped == sorted(values)
print("heap_pop: OK")
Hint 1

Save heap[0] as the result. Move the last element to position 0 (heap[0] = heap.pop()). Then sink: start at i = 0, find the smaller child, swap if needed, repeat.

Hint 2

Be careful with the edge case: if the heap has only one element, just return heap.pop(). When sinking, check that left_child(i) < len(heap) before accessing children. If only a left child exists (no right child), compare only with the left child.

Solution
def heap_pop(heap):
    if len(heap) == 1:
        return heap.pop()
    result = heap[0]
    heap[0] = heap.pop()
    i = 0
    while True:
        smallest = i
        l = left_child(i)
        r = right_child(i)
        if l < len(heap) and heap[l] < heap[smallest]:
            smallest = l
        if r < len(heap) and heap[r] < heap[smallest]:
            smallest = r
        if smallest == i:
            break
        heap[i], heap[smallest] = heap[smallest], heap[i]
        i = smallest
    return result

Exercise 3.4 — The Scheduler, Solved

Your heap works on anything comparable -- including tuples, which Python compares by first element. So a task is just (time, name).

Write run_scheduler(tasks) that takes a list of (time, name) tuples, pushes them all into a heap, then pops until empty, returning the task names in execution order.

run_scheduler([(1300, "task1"), (1230, "task2"), (1400, "task3"), (1100, "task4")])
# ["task4", "task2", "task1", "task3"]
Provided — tests
order = run_scheduler([(1300, "task1"), (1230, "task2"), (1400, "task3"), (1100, "task4")])
assert order == ["task4", "task2", "task1", "task3"]
print("run_scheduler: OK")
Hint 1

Push each task tuple into the heap. When you pop, you get (time, name). Collect the name part from each popped tuple into a result list.

Solution
def run_scheduler(tasks):
    h = []
    for task in tasks:
        heap_push(h, task)
    result = []
    while h:
        time, name = heap_pop(h)
        result.append(name)
    return result

Exercise 3.5 — Heap Sort, and the Standard Library

  1. You already proved it in the stress test: push n values, pop n values -> sorted. Package it as heap_sort(lst) returning a new sorted list. Each of the 2n operations costs O(log n) -- so heap sort is O(n log n), in place-able, no merging needed.
  2. Python ships your invention as the heapq module: heapq.heappush, heapq.heappop, and heapq.heapify (which builds a heap from a list in O(n)). Verify your heap agrees with it.
heap_sort([5, 2, 9, 1])   # [1, 2, 5, 9]
Provided — tests
import heapq

assert heap_sort([5, 2, 9, 1]) == [1, 2, 5, 9]
assert heap_sort([]) == []

# agree with the standard library on random data
data = [random.randint(0, 100) for _ in range(200)]
mine = list(data)
h = []
for v in mine:
    heap_push(h, v)

theirs = list(data)
heapq.heapify(theirs)

assert heap_pop(h) == heapq.heappop(theirs)
assert heap_pop(h) == heapq.heappop(theirs)
assert heap_pop(h) == heapq.heappop(theirs)
print("heap_sort + heapq agreement: OK")
Hint 1

Create an empty heap, push all values from lst, then pop them all into a result list. The result is sorted because the heap always gives you the minimum.

Solution
def heap_sort(lst):
    h = []
    for v in lst:
        heap_push(h, v)
    result = []
    while h:
        result.append(heap_pop(h))
    return result

Quick Check 3.6 — Heap Property

In a min-heap, which of the following is guaranteed?

Hint

The heap property is about parent-child relationships, not about siblings.

Reasoning

The min-heap property guarantees only that each parent <= its children. This means the root is the minimum (it's smaller than its children, which are smaller than theirs, etc.). But siblings have NO ordering guarantee -- the left child can be larger than the right child or vice versa. The heap is NOT fully sorted; it only guarantees the minimum is accessible in O(1) and can be removed in O(log n).

Quick Check 3.7 — Heap vs BST

Both a BST and a min-heap store elements. When would you choose a heap over a BST?

Hint

What operation is O(1) in a heap but O(log n) in a BST? And what operation is easy in a BST but hard in a heap?

Reasoning

A heap excels at one thing: quickly accessing and removing the min (or max). The root gives you the min in O(1), and removal is O(log n). But finding an arbitrary element requires scanning the entire heap -- O(n). A BST gives O(log n) search for any element and O(n) sorted traversal, but the minimum requires walking to the leftmost node -- O(log n). Choose a heap for priority queues (always need the "next most urgent"); choose a BST when you need general search or sorted iteration.

Part 4: Two More Tree Tricks

In-Order Without Recursion

Your in_order from Exercise 2.3 relies on Python's call stack. But recursion is not magic -- it's just a stack that Python manages for you. Prove it.

Write in_order_iter(node) that returns the same list as in_order, using no recursion -- manage a plain list as your own stack:

  1. From the current node, push nodes while walking left as far as you can.
  2. Pop a node -- that's the next value in order.
  3. Move to the popped node's right child and repeat, until both the stack and the current node are exhausted.
Provided — tests
assert in_order_iter(root) == in_order(root)
assert in_order_iter(None) == []

values = random.sample(range(1000), 100)
r = None
for v in values:
    r = insert(r, v)
assert in_order_iter(r) == sorted(values)
print("in_order_iter: OK -- recursion was just a hidden stack all along")
Hint 1

Initialize stack = [] and current = node. Loop while current is not None or stack. Inner loop: while current, push it and go left. Then pop, record the value, and set current = popped.right.

Hint 2
def in_order_iter(node):
    result, stack, current = [], [], node
    while current or stack:
        while current:
            stack.append(current)
            current = current.left
        current = stack.pop()
        result.append(current.val)
        current = current.right
    return result
Solution
def in_order_iter(node):
    result, stack, current = [], [], node
    while current or stack:
        while current:
            stack.append(current)
            current = current.left
        current = stack.pop()
        result.append(current.val)
        current = current.right
    return result

From Text to Tree: a Tiny Parser

In Part 1 you evaluated expression trees that were built by hand. Real programs receive expressions as strings: "((3+4)*(9-10))". Turning the string into a tree is called parsing -- it's the front half of every compiler (your evaluate from Exercise 1.1 is the back half).

Write parse(s) for fully parenthesized expressions with the operators +, -, * and non-negative integer numbers (no spaces):

A recursive helper that takes the current position i and returns (node, next_position) works beautifully.

Provided — tests
assert evaluate(parse("((3+4)*(9-10))")) == -7
assert evaluate(parse("(((3*4)+7)-12)")) == 7
assert evaluate(parse("(2+3)")) == 5

leaf = parse("42")
assert leaf.val == 42
assert leaf.left is None and leaf.right is None

print_tree(parse("((3+4)*(9-10))"))
print("parse: OK -- you just wrote the front half of a compiler")
Hint 1

Write a helper _parse(s, i) that returns (node, next_i). If s[i] is a digit, read all consecutive digits to form the number. If s[i] == '(', skip it, parse left, read the operator, parse right, skip ')'.

Hint 2

To read a multi-digit number: start = i; while i < len(s) and s[i].isdigit(): i += 1, then num = int(s[start:i]). The main parse(s) function calls _parse(s, 0) and returns just the node.

Solution
def parse(s):
    def _parse(s, i):
        if i < len(s) and s[i].isdigit():
            start = i
            while i < len(s) and s[i].isdigit():
                i += 1
            return Node(int(s[start:i])), i
        # must be '('
        i += 1  # skip '('
        left, i = _parse(s, i)
        op = s[i]
        i += 1  # skip operator
        right, i = _parse(s, i)
        i += 1  # skip ')'
        return Node(op, left, right), i
    node, _ = _parse(s, 0)
    return node
Structure You built Sweet spot
Expression tree evaluate -- recursion on trees how compilers see code
Binary search tree insert, contains, in_order, find_min, delete ordered data with frequent insert/delete, all O(log n) if balanced
Min-heap heap_push, heap_pop, run_scheduler, heap_sort only need the min: schedulers, priority queues, O(log n) with a plain list

What you just invented -- the trade-off ladder you climbed:

Choosing a data structure = paying only for the promises you actually use.

Bonus tricks: in_order_iter showed that recursion is just a stack you can manage yourself, and parse turned a string into an expression tree -- text in, evaluate-ready tree out.