Iterative Tree Traversals in Python

April 21, 2016

Tree traversals are most naturally expressed in recursion, but iterative versions are cool too, plus they take only O(1) space.

Inorder traversal: Visit the left subtree first, then the node, and the right subtree.

Preorder traversal: Visit the node first, then the left subtree, then the right subtree.

Postorder traversal: Visit the left subtree, then the right subtree, then the node.

The concept behind the iterative versions are as follows.

There are three states a traversal can be in:

  • You’ve just visited the left or right child of a parent node.
  • You’ve just gone back to a parent node from its left child.
  • You’ve just gone back to a parent node from its right child.

Keeping three pointers: prev to designate the previous node, curr to designate the current node, and _next to designate the next node, we can codify the above conditions like so:

if not prev or prev.left == curr or prev.right == curr:
    # first condition
elif curr.left == prev:
    # second condition
else: # curr.right == prev
    # third condition

With that in mind, I present the three different traversals, whose function signatures take a BTreeNode as the first argument and a function to operate on the tree nodes as the second argument.

class BTreeNode:
    def __init__(self, data):
        self.data = data
        self.parent = None
        self.left = None
        self.right = None

    def __str__(self):
        return str(self.data)

a = BTreeNode(6)
b = BTreeNode(4)
c = BTreeNode(5)
d = BTreeNode(3)
e = BTreeNode(2)
f = BTreeNode(1)

a.left = b
a.right = c

b.parent = a
b.left = d
b.right = e

c.parent = a
c.left = f

d.parent = b

e.parent = b

f.parent = c

def iterativeInOrder(root, func):
    if not root:
        return

    prev = None
    curr = root
    _next = None

    while curr:
        if not prev or prev.left == curr or prev.right == curr:
            if curr.left:
                _next = curr.left
            else:
                func(curr)
                _next = curr.right if curr.right else curr.parent

        elif curr.left == prev:
            func(curr)
            _next = curr.right if curr.right else curr.parent

        else:
            _next = curr.parent

        prev = curr
        curr = _next

def iterativePreOrder(root, func):
    if not root:
        return

    prev = None
    curr = root
    _next = None

    while curr:
        if not prev or prev.left == curr or prev.right == curr:
            func(curr)
            if curr.left:
                _next = curr.left
            else:
                _next = curr.right if curr.right else curr.parent

        elif curr.left == prev:
            _next = curr.right if curr.right else curr.parent

        else:
            _next = curr.parent

        prev = curr
        curr = _next

def iterativePostOrder(root, func):
    if not root:
        return

    prev = None
    curr = root
    _next = None

    while curr:
        if not prev or prev.left == curr or prev.right == curr:
            if curr.left:
                _next = curr.left
            elif curr.right:
                _next = curr.right
            else:
                func(curr)
                _next = curr.parent

        elif curr.left == prev:
            if curr.right:
                _next = curr.right
            else:
                func(curr)
                _next = curr.parent

        else:
            func(curr)
            _next = curr.parent

        prev = curr
        curr = _next

iterativeInOrder(a, print)   # 3 4 2 6 1 5
iterativePreOrder(a, print)  # 6 4 3 2 5 1
iterativePostOrder(a, print) # 3 2 4 1 5 6