Iterative Tree Traversal
By memorizing a simple implementation of iterative tree traversal we simplify a large number of programming interview questions.
Table of Contents
Introduction
One common programming interview problem asks candidates to write an in-order binary tree traversal without using recursion (LeetCode #94).
It's a deceptively simple problem, but I'll admit that I struggled the first time I saw it. It was a pen-and-paper interview, and I fumbled through it for a while before giving up and writing a recursive version, giving this excuse in the margin: “the recursive version is tidier - why wouldn’t you write it this way?”. Nice save[1].
Some of the online solutions to this are pretty unintuitive, so in this article I'll describe an approach that is easy to memorize or re-derive if needed. I've used this code verbatim in a surprising number of interview questions (some are in the Examples section). Here are some reasons you might want to do the same:
-
Iterators decouple traversal logic from the calculation logic. Factoring code into units of single responsibilities improves readability and reusability (useful in coding, but also the chunking technique used in rapid learning). Doing this in an interview signals good code hygiene, and can allow you to write less code (saving precious time)[2]. Admittedly, iterators don't have to be iterative, but by doing it iteratively...
-
You cover your bases. By making our iterators iterative we increase the number of questions we can mindlessly[3] apply it to. If you are asked to write a recursive version specifically, you can just do that instead (it should be easier anyway).
-
Avoiding recursion can prevent a stack overflow. Operating Systems limit stack space, and the typical recursive approach has non-optimizable tail calls, so we might face issues with large datasets. This is the main reason why you'd want to avoid recursion at all.
-
The wisdom of Earl Sweatshirt (Source):
Get your fundamentals on lock so that you can start getting into the ill advanced shit. This is universally applicable.
Let's get our fundamentals on lock!
Implementation
Recursive
No doubt you’ve seen recursive in-order tree traversal countless times, so let me instead give an example of an in-order iterator using Python 3.3's yield from
syntax:
def inorder(root):
if not root: return
yield from inorder(root.left)
yield root
yield from inorder(root.right)
Note: we yield root
instead of yield root.value
so that we can access the children if needed (like in the second example below).
By separating our traversal logic from the rest of our logic we end up with code that is easier to read and re-combine (see the examples). But we haven't answered the question - it is still recursive! If we implement it iteratively we can use the same code in more problems.
I recommend attempting to write an iterative version yourself before continuing (that link again: LeetCode #94). Pain (or slight discomfort) is the best teacher.
Iterative
All recursive functions can be written iteratively. In fact, like a compiler, we can mechanically convert our function calls to use an explicit stack, but readability/memorability will likely suffer.
While there are other iterative implementations (like this one using nested loops), I prefer the one below since it is less surprising[4], hence easier to remember or re-derive.
We start with a function to define the general approach:
def tree_iterator(root, expand):
frontier = [(False, root)]
while frontier:
expanded, curr = frontier.pop()
if not curr: continue
if expanded: yield curr
else: frontier.extend((x is curr, x) for x in reversed(expand(curr)))
This is just a Depth-First traversal[5] function, where an expand
function provides the order we would like to visit the children and parent, which are then reversed()
because stacks are LIFO (so nodes will be popped in reverse).
We don't use a visited set, since trees are acyclic, but we do mark the nodes as expanded or not so that we can re-visit each parent without expanding it twice.
We can then use tree_iterator
to define pre-order, in-order, and post-order:
preorder = lambda root: tree_iterator(root, lambda node: (node, node.left, node.right))
inorder = lambda root: tree_iterator(root, lambda node: (node.left, node, node.right))
postorder = lambda root: tree_iterator(root, lambda node: (node.left, node.right, node))
And their reversals:
reverse_preorder = lambda root: tree_iterator(root, lambda node: (node, node.right, node.left))
reverse_inorder = lambda root: tree_iterator(root, lambda node: (node.right, node, node.left))
reverse_postorder = lambda root: tree_iterator(root, lambda node: (node.right, node.left, node))
Each function is almost identical - all you have to do is remember the expansion order.
Examples
Here's a handful of example problems that are simplified by the above functions:
Count the number of nodes in a Binary Tree
sum(1 for _ in inorder(tree))
Count the number of leaf nodes in a Binary Tree
is_leaf = lambda x: not x.left and not x.right
sum(1 for x in inorder(root) if is_leaf(x)))
This is a good example of why our iterators yield the node instead of the value.
Flatten a Binary Search Tree into an Ordered List
[x.value for x in inorder(tree)]
Verify that a Binary Tree is a Binary Search Tree
from itertools import pairwise
is_ordered = lambda xs: all(a<=b for a,b in pairwise(xs))
is_ordered(x.value for x in inorder(tree))
Note: itertools
is probably one of the most useful modules in the Python standard library - get to know it intimately.
Find the kth-smallest element of Binary Search Tree
next(x.value for i,x in enumerate(inorder(root)) if i == k-1)
Find the kth-largest element of Binary Search Tree
next(x.value for i,x in enumerate(reverse_inorder(root)) if i == k-1)
Check if a Binary Tree is Symmetrical
from itertools import zip_longest
all(a.value == b.value for a,b in zip_longest(inorder(tree_a), reverse_inorder(tree_b)))
And of course the questions that ask for iterative versions of these, and the traversals themselves.
This is only a small sample of problems - if you find more examples please share them in the comments below!
Conclusion
Although the code is a little more complicated than the standard recursive approach, by implementing a non-recursive tree traversal iterator we can reuse the exact same code for a large amount of problems.
Try it out on some LeetCode tree questions for yourself and see if this approach lends itself to better solutions, and share any questions that benefit in the comments.
But don't stop there. Identifying patterns and compressing them into general solutions like this is an important part of deliberate practice. Get into the habit of repeating problems and solving them in new ways. Play Code Golf with yourself, and save the functions that you keep finding a use for[6]. Share those in the comments too!
Try to not be like I was back then. Intellectual Humility helps you learn faster, and makes you more enjoyable to work with :) ↩︎
This is true due to the potential for code reuse, but also, when I’m interviewing someone and they use a
helper(x)
function before implementing it, I sometimes ask them to skip implementing it in order to increase the signal-to-noise ratio. We only have an hour for an interview and need to be selective about what we assess. ↩︎Many people dislike the element of rote learning present in programming interviews today. While it is true that there should be better ways to gauge a mutual fit, the reality is that it's the current meta of the game we want to play, so is mostly unavoidable right now. However, repetition can help cache patterns so you can apply them more readily - in or out of an interview - so doing this isn't wasted time. ↩︎
The Principle of Least Astonishment is a UI design concept, but what is code if not a way for programmers to interface with both computers and other programmers? A fun hobbie is to go through UI/UX concepts and see how they apply to code. ↩︎
In-order, pre-order, and post-order are all Depth-First, whereas level-order is Breadth-First. ↩︎
I've started using Cacher to collect code snippets, but gists are great too. ↩︎
alexbowe.com Newsletter
Join the newsletter to receive the latest updates in your inbox.