Skip to content

Iterative Tree Traversal

By memorizing a simple implementation of iterative tree traversal we simplify a large number of programming interview questions.

Alex Bowe
Alex Bowe
5 min read
An aerial view of a forest.

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:

  1. 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...

  2. 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).

  3. 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.

  4. 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!


  1. Try to not be like I was back then. Intellectual Humility helps you learn faster, and makes you more enjoyable to work with :) ↩︎

  2. 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. ↩︎

  3. 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. ↩︎

  4. 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. ↩︎

  5. In-order, pre-order, and post-order are all Depth-First, whereas level-order is Breadth-First. ↩︎

  6. I've started using Cacher to collect code snippets, but gists are great too. ↩︎

pythonalgorithmstreesprogramming interviewslearninggraph searchfundamentals

Alex Bowe Twitter

Alex has a PhD in succinct data structures for bioinformatics, and currently works on self-driving cars. He's also interested in cryptocurrency, business, fashion, photography, teaching, and writing.


Related Articles

Members Public

Newsletter Zero

My first newsletter, where I cover how I want to write more, how I finished my PhD, what I'm working on at Cruise, recent articles I've written, a programming interview course I'm working on, and previews for future posts.

Members Public

How to Recover a Bitcoin Passphrase

How I recovered a Bitcoin passphrase by performing a Breadth-First search on typos of increasing Damerau-Levenshtein distances from an initial guess.

How to Recover a Bitcoin Passphrase
Members Public

Succinct de Bruijn Graphs

This post will give a brief explanation of a Succinct implementation for storing de Bruijn graphs [http://en.wikipedia.org/wiki/De_Bruijn_graph], which is recent (and continuing) work I have been doing with Sadakane. Using our new structure, we have squeezed a graph for a human genome (which