Iterative Tree Traversals in Python

April 21, 2016

Tree tra­ver­sals are most nat­u­ral­ly expressed in recur­sion, but iter­a­tive ver­sions are cool too, plus they take only O(1) space.

Inorder tra­ver­sal: Visit the left sub­tree first, then the node, and the right sub­tree.

Pre­order tra­ver­sal: Visit the node first, then the left sub­tree, then the right sub­tree.

Pos­torder tra­ver­sal: Visit the left sub­tree, then the right sub­tree, then the node.

The con­cept behind the iter­a­tive ver­sions are as fol­lows.

There are three states a tra­ver­sal can be in:

  • You’ve just vis­it­ed 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.

Keep­ing three point­ers: prev to des­ig­nate the pre­vi­ous node, curr to des­ig­nate the cur­rent node, and _next to des­ig­nate the next node, we can codify the above con­di­tions 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 pre­sent the three dif­fer­ent tra­ver­sals, whose func­tion sig­na­tures take a BTreeNode as the first argu­ment and a func­tion to oper­ate on the tree nodes as the second argu­ment.

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