Tree traversal made hard
Tree traversal is a well known task for which recursion is simple and natural. We define a Node class.
class Node:
def __init__(self, value, left=None, right=None):
self.value = value
self.left = left
self.right = right
def __repr__(self):
return repr(self.value)
@classmethod
def make_tree(cls, values):
if len(values) == 0:
return None
mid = len(values) // 2
return Node(
values[mid],
left=cls.make_tree(values[:mid]),
right=cls.make_tree(values[mid + 1 :]),
)
The class method make_tree returns the root node of a balanced binary tree from a list of values. If the list is sorted, then the tree is a binary search tree (BST). A BST has the important property that if you do an in-order traversal of the nodes, you will visit them in order. A recursive implementation of in-order traversal is straightforward.
def traverse_recursive(node):
if node.left:
traverse_recursive(node.left)
print(node)
if node.right:
traverse_recursive(node.right)
Let’s test it with a small sorted list of integers and verify that the nodes are printed out in order.
>>> tree = Node.make_tree([1, 2, 3, 4, 5, 6])
>>> traverse_recursive(tree)
1
2
3
4
5
6
Now, what if we insist on traversing without recursion? In lieu of the call stack, we need to maintain our own stack. Each item in our stack would hold the same information as a call frame. In traverse_recursive there are two potential recursive calls, each call adding another frame to the call stack. Since a recursive call suspends the calling frame, our stack must maintain a state - whether or not the left side of the BST has been examined. Let’s designate the possibilities for this state "left" and "right".
def traverse_iterative(node):
stack = []
stack.append((node, "right"))
stack.append((node, "left"))
while stack:
node, state = stack.pop()
if state == "left":
if node.left:
stack.append((node.left, "right"))
stack.append((node.left, "left"))
elif state == "right":
print(node)
if node.right:
stack.append((node.right, "right"))
stack.append((node.right, "left"))
Each call frame becomes, in our stack, a pair of items - as the stack is popped off, the left side of each node is examined before the right. The analogue of a recursive call is pushing another pair to the stack. Let’s test traverse_iterative on the same BST.
>>> traverse_iterative(tree)
1
2
3
4
5
6
Note that an in-order traversal is different from a depth-first search (DFS). To implement DFS iteratively, we simply maintain a stack of nodes without the additional state. This is because DFS first checks whether the target node has been found (and if so, exits) before adding the child nodes to the stack. Hence there is no reason to “remember” the parent node. Formally, DFS on a binary tree is equivalent to a pre-order traversal.
The iterative implementation of in-order tree traversal is less intuitive than the recursive version, which nicely illustrates that recursion is sometimes much easier than iteration. Though iteration has the advantage of not growing the call stack, and readily leads to tail recursion.
def traverse_tail_recursive(stack):
if not stack:
return
node, state = stack.pop()
if state == "left":
if node.left:
stack.append((node.left, "right"))
stack.append((node.left, "left"))
elif state == "right":
print(node)
if node.right:
stack.append((node.right, "right"))
stack.append((node.right, "left"))
traverse_tail_recursive(stack)
Calling traverse_tail_recursive with the same tree but wrapped inside of a stack:
>>> traverse_tail_recursive(
... [
... (tree, "right"),
... (tree, "left")
... ]
... )
1
2
3
4
5
6
Python does not offer tail call optimization, but in languages that do, this example suggests a recipe for turning a recursive function into a tail recursive function: via an intermediate iterative implementation.