Iterative traversals with an explicit stack in Python

Why this matters

In the first lesson you visited the nodes of a tree with recursion. A recursive function calls itself for the left child and then for the right child, and Python keeps track of where each call should return. That bookkeeping is hidden. You never see it.

Some situations need the bookkeeping in your own hands. A tree can be deep enough to make recursion run out of room, which Python reports as a RecursionError once the call chain passes its limit of 1000 calls. An interview question may say "do it without recursion". A program may need to pause a traversal halfway and resume it later. In all three cases you write the traversal with a loop and a stack that you control.

This lesson covers preorder and inorder. The next lesson covers postorder, which is harder because a parent has to wait for its children. You will use the same tree as the first lesson, so you can compare the two versions directly.

The tree we will use

Here is the tree with six nodes. The root is 1. Node 2 is the left child of 1, node 3 is the right child of 1, node 4 and node 5 are the children of 2, and node 6 is the right child of 3.

        1
       / \
      2   3
     / \   \
    4   5   6

In Python, a node is an object with a val attribute and two attributes, left and right. An attribute holds a reference to another object, or None when there is no child. The class and the builder function look like this:

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
 
 
def build_sample_tree():
    n1 = TreeNode(1)
    n2 = TreeNode(2)
    n3 = TreeNode(3)
    n4 = TreeNode(4)
    n5 = TreeNode(5)
    n6 = TreeNode(6)
    n1.left = n2
    n1.right = n3
    n2.left = n4
    n2.right = n5
    n3.right = n6
    return n1

Every example below starts from build_sample_tree(). The lessons in this module use val for the value of a node, as the earlier lessons did.

The stack, in plain words

A stack is a pile. You can put one thing on the top (this is a push) and you can take the top thing off (this is a pop). You can never reach the thing under the top until the things above it are gone. Computer scientists call this rule last in, first out, or LIFO.

A pile of plates is a good picture. The plate you put down last is the first one you pick up.

Python does not have a separate stack type. A plain list does the job, and it is the usual choice. The methods you need are these:

• stack.append(x) puts x on the top. The top is the right end of the list.
• stack.pop() removes the top and returns it. It raises an error if the list is empty.
• stack[-1] returns the top without removing it. It raises an error if the list is empty, so check first.
• not stack is True when the stack has nothing in it.

Both append and pop work at the right end of the list, so each one takes a short, fixed time no matter how long the list is.

The code below shows the basic methods on their own, with the stack printed after each step. Python prints a list starting from the left end, which is the bottom of the pile. To print the stack from the top, as the traces in this lesson do, the code reverses the list before it prints it.

Output:

after three pushes: [30, 20, 10]
peek: 30
pop: 30
after one pop: [20, 10]

The top of the stack is the first value in the printed list, and every example in this lesson prints the stack the same way.

Example 1: preorder with a stack

Preorder means: visit a node, then visit its left subtree, then visit its right subtree. The node comes before its children. For the sample tree the answer is 1 2 4 5 3 6.

With recursion, the order falls out of the call structure. With a stack, you have to arrange the pushes so that the pops come out in the same order. The rule is simple. When you pop a node and visit it, push its right child first and its left child second. The left child is then on top, so it is popped next. That gives the left subtree a full turn before the right child is touched.

Here is the function. It prints a line after every push and every pop, so you can see the stack change:

The helper show prints the values in the stack from top to bottom. It uses reversed to start from the right end of the list, which is the top. The join call puts , between the values, and the first and last brackets come from the string itself, so the output has no extra spaces.

The output of this program is the trace for the sample tree. Read it one line at a time:

push 1 -> stack [1]
pop  1 -> visit, stack []
push 3 -> stack [3]
push 2 -> stack [2, 3]
pop  2 -> visit, stack [3]
push 5 -> stack [5, 3]
push 4 -> stack [4, 5, 3]
pop  4 -> visit, stack [5, 3]
pop  5 -> visit, stack [3]
pop  3 -> visit, stack []
push 6 -> stack [6]
pop  6 -> visit, stack []
result: [1, 2, 4, 5, 3, 6]

The first push puts the root on the stack. The first pop takes it off and visits it, so 1 is the first value in the result. Then node 1 has both children. The right child 3 goes on first, then the left child 2 goes on top of it. The next pop takes 2, which is the left child, so the left side is handled first.

Look at what happens to node 2. Its right child 5 is pushed, then its left child 4 is pushed on top. Node 4 is a leaf, so popping it adds nothing new. The next pop is 5, the right child of 2. Only after both children of 2 are done does the stack show 3 on top, and then node 3 is visited. Node 3 pushes its right child 6, which is the last node.

The result is [1, 2, 4, 5, 3, 6]. That matches the preorder definition: each node comes before its children, and the left subtree comes before the right subtree.

Why the order of the two pushes matters

Suppose you swapped the two pushes so that the left child went on first. The left child would sit below the right child, so the right child would be popped first. The traversal would then visit the right side before the left side. The answer would still look like a traversal, but it would be the mirror image of the one you want. A later lesson on mirror trees uses exactly this idea on purpose.

Example 2: inorder with a stack

Inorder means: visit the left subtree, then the node itself, then the right subtree. For the sample tree the answer is 4 2 5 1 3 6. The value of each node comes between the values of its left subtree and its right subtree.

Inorder is harder to do with a single push-and-pop rule, because the node is visited in the middle. You cannot visit a node as soon as you see it. You first have to go as far left as possible, and only then can you visit the nodes on the way back up.

The function uses two things. One is the stack, which holds nodes you have seen but not yet visited. The other is a variable called current, which holds the node you are walking down from. The loop does two jobs. First it walks left, pushing every node on the way. Then it pops a node, visits it, and moves to its right child.

def inorder(root):
    visited = []
    stack = []
    current = root
    while current is not None or stack:
        while current is not None:
            stack.append(current)
            print(f"push {current.val} -> stack {show(stack)}")
            current = current.left
        node = stack.pop()
        print(f"pop  {node.val} -> visit, stack {show(stack)}")
        visited.append(node.val)
        current = node.right
    return visited

The outer while runs as long as there is something to do. That is either a node we are walking down from (current is not None) or a node waiting on the stack. The inner while is the walk to the left. It pushes current, moves to its left child, and keeps going until the left child is None.

The function uses the show helper and the TreeNode class from the preorder example, so put them in the same file. Here is the trace on the sample tree:

push 1 -> stack [1]
push 2 -> stack [2, 1]
push 4 -> stack [4, 2, 1]
pop  4 -> visit, stack [2, 1]
pop  2 -> visit, stack [1]
push 5 -> stack [5, 1]
pop  5 -> visit, stack [1]
pop  1 -> visit, stack []
push 3 -> stack [3]
pop  3 -> visit, stack []
push 6 -> stack [6]
pop  6 -> visit, stack []
result: [4, 2, 5, 1, 3, 6]

The first three pushes walk down the left side: 1, then 2, then 4. Node 4 has no left child, so the walk stops. The first pop takes 4 off and visits it. Node 4 has no right child, so current becomes None, and the inner while does nothing. The next pop takes 2 off and visits it. Node 2 has a right child, so current is set to 5. The inner while pushes 5, finds no left child, and the next pop visits 5.

The stack is now [1]. Node 1 is popped and visited, and its right child 3 becomes current. Node 3 has no left child, so it is pushed and immediately popped. Then its right child 6 is pushed and popped. The stack is empty and current is None, so the outer loop ends.

Notice that the visit happens at the moment of the pop. That is the step that puts the node in the middle of the output. A node is pushed when you walk past it on the way left, and it is popped and visited only after its whole left side is finished.

Tracing the stack size

In the sample tree the largest stack in the trace has three nodes: [4, 2, 1]. Those are the nodes on the path from 1 down to 2 and then to 4.

Example 3: a long left chain and an empty tree

A tree where every node has only a left child is called a left chain. It is the worst case for a stack-based inorder, because the walk to the left goes all the way down before the first visit. Here is a chain of five nodes, built by a helper that you can add to the same file:

def build_left_chain(count):
    root = TreeNode(1)
    current = root
    for v in range(2, count + 1):
        current.left = TreeNode(v)
        current = current.left
    return root

The trace of inorder(build_left_chain(5)) looks like this:

push 1 -> stack [1]
push 2 -> stack [2, 1]
push 3 -> stack [3, 2, 1]
push 4 -> stack [4, 3, 2, 1]
push 5 -> stack [5, 4, 3, 2, 1]
pop  5 -> visit, stack [4, 3, 2, 1]
pop  4 -> visit, stack [3, 2, 1]
pop  3 -> visit, stack [2, 1]
pop  2 -> visit, stack [1]
pop  1 -> visit, stack []
result: [5, 4, 3, 2, 1]

All five nodes are pushed before the first pop, so the stack holds five items at its largest. A chain of one million nodes would hold one million items, so the stack grows with the tree. You should know that size before you use this method on a large input.

The output is the values in reverse order. A left chain has no right children, so inorder visits each node after everything below it, and that reverses the order from top to bottom.

The empty tree is the last case. If root is None, the outer while in inorder never runs, because current is None and the stack is empty. The result is an empty list. The preorder function checks for None at the top, which is a clear way to show the reader the same thing.

preorder: []
inorder:  []

A program that prints both results for None shows these two lines.

Common mistakes

Popping from an empty stack

pop() on an empty list raises an IndexError. The most common way to reach this is to leave out the loop test current is not None or stack, so the loop never stops on its own. The example below has that mistake. The pop at the bottom then runs when nothing is left. The code is shown without the run mark because it stops with an error.

def inorder_no_guard(root):
    visited = []
    stack = []
    current = root
    while True:
        while current is not None:
            stack.append(current)
            current = current.left
        node = stack.pop()
        visited.append(node.val)
        current = node.right
 
 
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
 
 
root = TreeNode(1)
root.left = TreeNode(2)
print(inorder_no_guard(root))

Running it on a root with a left child, a tree of two nodes, produces this real output:

Traceback (most recent call last):
  File "main.py", line 23, in <module>
    print(inorder_no_guard(root))
          ~~~~~~~~~~~~~~~~^^^^^^
  File "main.py", line 9, in inorder_no_guard
    node = stack.pop()
IndexError: pop from empty list

The last line says what happened: the program asked the stack for an item and there was none. The line above the error shows the call that failed, which is the stack.pop() call inside inorder_no_guard. The fix is the loop condition current is not None or stack, which is the one used in the correct version.

Pushing the children in the wrong order in preorder

If you push the left child first, the right child is popped first. The output is then the mirror image of preorder. No error is raised, so this mistake is easy to miss. Check it with a small tree whose answer you know. The tree in the exercises below shows the difference clearly.

Setting current to the wrong node

In inorder, the step after a visit must be current = node.right. If you write current = node.left, the walk goes back down the left side that you already finished, and nodes can be visited twice. If you forget the line altogether, current stays None, so the loop only pops what is left on the stack. For the sample tree the result is [4, 2, 1], and the right subtrees are missing.

Confusing pop and peek

stack[-1] looks at the top without removing it. A loop that reads stack[-1] where it should call pop() never removes the node, so it keeps seeing the same node. The inorder function above calls pop() on purpose, because each node must leave the stack when it is visited. The next lesson uses stack[-1] on purpose, for a reason you will see there.

Forgetting the None check on the root

If root is None, code that starts with stack.append(root) and then reads root.val raises an AttributeError. The preorder function returns early for that reason. Test the empty tree every time you write a traversal. It is a common input in interviews and in real code. The code is shown without the run mark because it stops with an error.

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
 
 
def preorder_no_check(root):
    stack = [root]
    visited = []
    while stack:
        node = stack.pop()
        visited.append(node.val)
        if node.right is not None:
            stack.append(node.right)
        if node.left is not None:
            stack.append(node.left)
    return visited
 
 
print(preorder_no_check(None))

The output is:

Traceback (most recent call last):
  File "main.py", line 21, in <module>
    print(preorder_no_check(None))
          ~~~~~~~~~~~~~~~~~^^^^^^
  File "main.py", line 13, in preorder_no_check
    visited.append(node.val)
                   ^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'val'

The first pop takes None off the stack, and the next line reads None.val. The fix is to return early when root is None, as the preorder function does.

Practice

Exercise 1

Here is a small tree with four nodes. The root is 10. Its left child is 5, and 5 has a left child 2. The root has a right child 15. Write down the preorder and the inorder of this tree. Then trace the stack for the preorder function.

Show the answer

Preorder is 10 5 2 15. Inorder is 2 5 10 15. For preorder, the trace is: push 10, pop 10 and visit, push 15, push 5, pop 5 and visit, push 2, pop 2 and visit, pop 15 and visit. Each pop is one visit, and the stack is empty at the end.

Exercise 2

Add a node 7 as the right child of node 4 in the sample tree. Node 4 already has no children, so 7 becomes its only child. What is the inorder traversal of the new tree? Work it out by hand, then check it by running the code.

Show the answer

The inorder traversal is 4 7 2 5 1 3 6. Node 4 has the right child 7, so 7 comes right after 4. The left subtree of 1 is 4 7 2 5, then comes 1, then the right subtree 3 6.

Exercise 3

Preorder pushes the right child first and the left child second. Explain in your own words why the left child is visited first. Then say what changes if you swap the two pushes, using the tree from Exercise 2.

Show the answer

A stack is last in, first out. The child pushed last is on top, so it is popped first. The left child is pushed last, so it is visited first. If you swap the pushes, the right child is on top and is visited first. For the tree in Exercise 2, the swapped version gives 1 3 6 2 5 4 7 instead of 1 2 4 7 5 3 6.

Exercise 4

Write a function count_nodes(root) that counts the nodes without recursion. Use the inorder loop as your starting point. Each node is visited exactly once, so counting the visits is enough.

Show the answer

def count_nodes(root):
    count = 0
    stack = []
    current = root
    while current is not None or stack:
        while current is not None:
            stack.append(current)
            current = current.left
        node = stack.pop()
        count += 1
        current = node.right
    return count

Add count = 0 before the loop, add count += 1 right after the pop, and return count at the end. The function returns 7 for the tree in Exercise 2, 4 for the four-node tree in Exercise 1, and 0 for None. A simple way to check the count is to compare it with the length of the list that inorder returns.

Exercise 5

A tree is built as a left chain of 1000 nodes. Using the inorder function from this lesson, how many items can the stack hold at its largest? Explain your answer in one sentence.

Show the answer

It holds 1000 items. Every node is on the left path, so the walk pushes all of them before the first pop. The stack is at its largest at that moment.

Summary

• A stack is a last in, first out pile. In Python, a list with append, pop and stack[-1] is the stack.
• For preorder, push the right child before the left child, so the left child is popped and visited first.
• For inorder, walk left and push every node on the way, then pop, visit and move to the right child.
• The inorder loop stops when current is None and the stack is empty. Leave out either part and you get a wrong answer or an error.
• The stack holds the nodes you have seen but not yet visited, so a left chain of five nodes fills it to five items.