BST iterator with a stack in Python

Why this matters

Sometimes you do not need every value in a tree at once. You may want the smallest value, then the next one, and stop as soon as you have enough. You may also want to read two sorted trees side by side and merge them into one sorted sequence.

The simple way is to build a full list of all the values and then read the list. That works. But the list holds every value in memory before you read the first one.

An iterator avoids that. An iterator is an object that gives you one value per call. In Python, a class becomes an iterator by having a __next__ method, which returns the next value, and by raising StopIteration when no values are left. The built-in function next(it) calls __next__ for you. A for loop does the same thing until StopIteration appears.

This lesson builds an iterator for a binary search tree (BST). It keeps a stack, which is a pile where you only add and remove at the top. In Python, a plain list works as a stack: append adds to the top and pop removes from the top. The stack holds only the nodes the walk still needs. Its size depends on the height of the tree, not on how many values the tree holds.

Words used in this lesson

• A node is one box in the tree. It holds one value and up to two links, to a left child and a right child.
• A child is a node directly below another node. A parent is the node above it.
• A subtree is a node together with everything below it.
• A BST is a tree where every value in a node's left subtree is smaller than the node, and every value in its right subtree is larger.
• Inorder means: visit the left subtree, then the node, then the right subtree.
• The height is the number of levels. A tree with one node has height 1.
• A stack is a last-in, first-out pile. append adds to the top, pop removes from the top, and stack[-1] looks at the top without removing it.

Inorder order with a small example

Insert 7, 3, 15, 1, 5, 12, 20 in that order. Each new value goes left of a node if it is smaller and right if it is larger. The tree looks like this:

          7
        /   \
       3     15
      / \    / \
     1   5  12  20

Read the tree in inorder. Start with the left subtree of 7, which is the subtree rooted at 3. Its inorder is 1, 3, 5. Then comes 7. Then the right subtree rooted at 15 gives 12, 15, 20. The full result is:

1, 3, 5, 7, 12, 15, 20

Inorder on a BST always gives the values in sorted order. The iterator in this lesson does the same walk, but it stops after each value and waits for the next call.

These helper functions build the tree in the code examples. The examples below use this TreeNode class and the two functions next to it.

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
 
 
def insert(root, val):
    if root is None:
        return TreeNode(val)
    if val < root.val:
        root.left = insert(root.left, val)
    elif val > root.val:
        root.right = insert(root.right, val)
    return root
 
 
def build(*values):
    root = None
    for v in values:
        root = insert(root, v)
    return root

Two details matter here. First, insert returns the root of the tree it was given. When the tree is empty, the root is None, so insert creates a new node and returns it. The caller must store that result, which is what build does with root = insert(root, v). Second, *values in the definition of build is a variable-length parameter. It means the caller can pass any number of values, and Python packs them into a tuple. The call build(7, 3, 15) is the same as passing one tuple of three values.

The full-list approach first

The easy version collects the values in a list with a recursive inorder walk.

def fill(node, out):
    if node is None:
        return
    fill(node.left, out)
    out.append(node.val)
    fill(node.right, out)
 
 
def inorder_list(root):
    out = []
    fill(root, out)
    return out

Calling inorder_list(build(7, 3, 15, 1, 5, 12, 20)) returns:

[1, 3, 5, 7, 12, 15, 20]

Let n be the number of nodes. Each node is visited once, so the time is O(n). O(n) means the work grows in step with n: twice as many nodes takes about twice as long. The list also holds n values, so the extra memory is O(n) as well.

For a tree with 1,000,000 values, the list holds 1,000,000 entries before you read the first one. That is fine when you need all of them. It is wasteful when you only need the first few.

The idea: a pile of nodes still owed

Think of the walk as a pile of index cards. Each card is a node whose value you still owe to the caller. The card on top is the next value to return.

At the start, the smallest value is the far-left node. To reach it from the root, you go left again and again. Each node on that trip is a value you will return later, so each one goes on the pile. The helper that does this is called _push_left_path:

_push_left_path(node):
    while node is not None:
        push node onto the pile
        move node to its left child

When the caller asks for the next value, the iterator takes the top card (pops it), returns its value, and then handles the right child. Why the right child? Its values are all larger than the popped node. They are also smaller than the next card under it, because the tree is sorted. So the right child's left path goes on the pile next. If the right child is None, nothing is pushed.

That is the whole method. Two rules: start with the left path from the root, and after each pop, push the left path of the popped node's right child.

Worked example 1: the small tree

Here is the iterator class. It is the code behind the stack traces and the worked examples.

class BstIterator:
    def __init__(self, root):
        self.stack = []
        self.max_size = 0
        self._push_left_path(root)
 
    def _push_left_path(self, node):
        while node is not None:
            self.stack.append(node)
            node = node.left
        self.max_size = max(self.max_size, len(self.stack))
 
    def __iter__(self):
        return self
 
    def __next__(self):
        if not self.stack:
            raise StopIteration
        node = self.stack.pop()
        self._push_left_path(node.right)
        return node.val
 
    def peek(self):
        if not self.stack:
            raise IndexError("no more values")
        return self.stack[-1].val
 
    def stack_values(self):
        return [n.val for n in reversed(self.stack)]

The method names that start with one underscore, such as _push_left_path, mark a helper that the caller should not use directly. The method __iter__ returns the object itself, so the class works in a for loop. The method __next__ does the work. stack_values returns the stack from top to bottom as a list, which is what the trace below shows. The list's last item is the top of the stack, so reversed puts the top first.

The main part of the program calls next(it) until StopIteration appears:

root = build(7, 3, 15, 1, 5, 12, 20)
it = BstIterator(root)
print("stack after construction (top first):", it.stack_values())
while True:
    try:
        v = next(it)
    except StopIteration:
        break
    print(f"next() = {v}   stack now (top first): {it.stack_values()}")

The real output is:

stack after construction (top first): [1, 3, 7]
next() = 1   stack now (top first): [3, 7]
next() = 3   stack now (top first): [5, 7]
next() = 5   stack now (top first): [7]
next() = 7   stack now (top first): [12, 15]
next() = 12   stack now (top first): [15]
next() = 15   stack now (top first): [20]
next() = 20   stack now (top first): []

Read the stack column in the table form, so each step is easy to follow:

Three steps show the rule at work.

The constructor starts at 7, pushes 7, then 3, then 1, because each is the left child of the one before. Now the top is 1, the smallest value.

The first next() pops 1. It has no right child, so nothing is pushed. The top is now 3.

The second next() pops 3. Its right child is 5. The code pushes the left path of 5, which is just 5 itself. The stack is now 5 on top of 7.

The fourth next() pops 7. Its right child is 15. The code pushes the left path of 15: first 15, then its left child 12 on top. So the stack is 12 on top of 15. That is why the value after 7 is 12, and 15 waits below it.

When next raises StopIteration, the stack is empty and the walk is finished.

Why the stack stays small

Each node on the stack is an ancestor of the node above it. In the table, 3 is a child of 7, 1 is a child of 3, and so on. So all the nodes on the stack lie on one path that goes down from the root. A path has at most one node on each level, so the stack never holds more nodes than the tree has levels.

The small tree has three levels: 7 alone, then 3 and 15, then 1, 5, 12 and 20. Its height is 3, and the largest stack in the walk was 3 as well.

The iterator therefore uses O(h) memory, where h is the height. The memory depends on the shape of the tree, not on how many values it contains.

Worked example 2: the same values in three shapes

The stack size depends on shape. This example builds three trees with the same seven values and records the largest stack size during the walk. The class keeps the value in max_size, and the for loop reads every value.

shapes = [
    [7, 3, 15, 1, 5, 12, 20],
    [1, 2, 3, 4, 5, 6, 7],
    [7, 6, 5, 4, 3, 2, 1],
]
names = ["balanced", "increasing inserts", "decreasing inserts"]
for name, vals in zip(names, shapes):
    b = BstIterator(build(*vals))
    count = 0
    for _ in b:
        count += 1
    print(f"{name}: values={count} max stack={b.max_size}")

Output:

balanced: values=7 max stack=3
increasing inserts: values=7 max stack=1
decreasing inserts: values=7 max stack=7

The increasing inserts give a tree that leans right. Every node has only a right child, so the constructor pushes just the root, and each next() pushes one right child at most. The stack never holds more than one node.

The decreasing inserts give a tree that leans left. The constructor pushes all seven nodes down the left side, so the stack holds seven nodes at the start. The height is 7, and the stack reached 7.

The balanced tree has height 3, and the stack reached 3. In every case the largest stack size equals the height here, and it can never be larger than the height.

The full-list approach would hold 7 values in all three cases. A balanced tree with a million values has height about 20, so its stack would hold about 20 entries, while the list would hold a million.

Worked example 3: merge two BSTs in sorted order

Two iterators let you walk two trees at the same time. This is the merge step from merge sort, applied to two sorted sequences that come from BSTs. The code keeps one current value from each iterator. It outputs the smaller one and moves that iterator forward. The built-in next(it, None) returns None when an iterator is finished, instead of raising StopIteration. The second argument is the value to return in that case.

a = build(2, 6, 10)
b = build(1, 4, 6, 9)
ia = BstIterator(a)
ib = BstIterator(b)
merged = []
x = next(ia, None)
y = next(ib, None)
while x is not None or y is not None:
    if y is None or (x is not None and x <= y):
        merged.append(x)
        x = next(ia, None)
    else:
        merged.append(y)
        y = next(ib, None)
print("merged:", merged)

Output:

merged: [1, 2, 4, 6, 6, 9, 10]

Look at the if condition. If y is None, the iterator for b is finished, so the code takes x. The check y is None comes first, so the code never compares a missing value. If y is not None, the code only compares x <= y after confirming x is not None. The loop only runs while at least one of them is not None, so x is never None when the first branch adds it.

Both trees contain a 6. Both copies appear in the output, in the order the trees gave them.

The merged list is the only large structure, and it holds every value in the output. The two iterators add at most the two heights in stack entries, and neither tree is copied.

Common mistakes

Calling next() when nothing is left

If the code calls next(it) after the last value, the iterator raises StopIteration. On an iterator built from an empty tree, the first call fails:

it = BstIterator(None)
next(it)

The output ends with a traceback. Its last two lines are:

    raise StopIteration
StopIteration

StopIteration has no message of its own. It is the signal that the sequence has ended, so a for loop catches it for you, and a direct call needs a try block or the default value of next(it, None), as the merge example does.

Pushing the right child but not its left path

A common wrong version pushes the right child alone:

def buggy_walk(root):
    out = []
    stack = []
    node = root
    while node is not None:
        stack.append(node)
        node = node.left
    while stack:
        n = stack.pop()
        out.append(n.val)
        if n.right is not None:
            stack.append(n.right)
    return out

This looks plausible, but it skips values. The node 12 is the left child of 15, so 12 is reachable only through the left path of 15. If that path is never pushed, 12 never appears. Running this version on the small tree prints:

buggy walk: [1, 3, 5, 7, 15, 20]

The value 12 is missing. The fix is to call self._push_left_path(node.right), which pushes the whole left chain.

Using peek() where pop() is needed

If __next__ reads the top node with self.stack[-1] and never removes it, the next call sees the same node again. For the leaf 1 in the small tree, the stack never changes, so the loop never ends. The fix is pop(), which removes the node from the stack.

Building the stack from the right side

If the constructor pushes the right path instead of the left path, the walk goes wrong. A BST iterator for ascending values must push the left path in the constructor and in __next__. Pushing the right path first would return the largest value first, and the walk would print the values in decreasing order.

Practice

Exercise 1: trace the stack

Build a tree by inserting 4, 2, 6, 1, 3, 5, 7 in that order. Write the stack, top first, after the constructor and after each of the first three next() calls.

Show the answer

After the constructor the stack is 1, 2, 4. The first next() returns 1 and leaves 2, 4. The second returns 2, pushes 3 (the left path of its right child), and leaves 3, 4. The third returns 3 and leaves 4. Running stack_values() after each call prints the same lists.

Exercise 2: add peek()

Write the peek method so that it returns the next value without removing it. It must raise an error when the iterator is empty.

Show the answer

def peek(self):
    if not self.stack:
        raise IndexError("no more values")
    return self.stack[-1].val

The top of the stack is always the next value, so peek reads it without changing anything. On the tree built from 5, 3, 8, 1, 4, 7, 9, peek returns 1, and after one next() it returns 3.

Exercise 3: maximum stack size for an increasing insert

Build a tree by inserting 1, 2, 3, ..., 15 in that order. What is the largest stack size during a full walk?

Show the answer

The answer is 1. Each insert goes to the right of the last one, so the tree is a straight chain of right children.

Exercise 4: why 4 and not 15

Build a full tree with 15 nodes by inserting 8, 4, 12, 2, 6, 10, 14, 1, 3, 5, 7, 9, 11, 13, 15 in that order. The tree has 4 levels. Explain why the largest stack size is 4, not 15.

Show the answer

The stack holds ancestors that lie on one path from the root down, and a path has at most one node per level. This tree has 4 levels, so the stack never holds more than 4 nodes.

Exercise 5: the first k values

Write a function first_k(root, k) that returns the first k values in sorted order. It must stop early, so it does not walk the rest of the tree.

Show the answer

def first_k(root, k):
    out = []
    it = BstIterator(root)
    while len(out) < k:
        try:
            out.append(next(it))
        except StopIteration:
            break
    return out

For the tree built from 5, 3, 8, 1, 4, 7, 9, the function returns [1, 3, 4] for k equal to 3, [1, 3, 4, 5, 7, 8, 9] for k equal to 100, and [] for k equal to 0. The except branch stops the loop when the values run out, so asking for more values than exist is safe.

Summary

• An iterator gives one value per call, so the caller decides when to stop.
• The walk keeps a stack of nodes still owed. The top of the stack is the next value in sorted order.
• Each __next__ call pops one node and pushes the left path of its right child.
• The stack holds one node per level at most, so its size is O(h), the height of the tree.
• A full list of values costs O(n) memory, while the iterator's stack costs memory that depends on the height.