# Using generators for on-demand tree traversal

Let’s take a short break from the ongoing series of posts on AI to look at an interesting problem and a useful accompanying concept.

Problem: Given a BST and a value k, find two elements in the BST which sum to k.

If we were solving this problem for a sorted array, we could use the following two pointer approach to solve it in linear time:

Things get more interesting when we have to work with a BST. An obvious solution is to store the inorder traversal of the BST in an array and now, we’ve reduced the problem to finding two elements in a sorted array which sum to k. This takes linear space and time, which would be optimal if the BST was unbalanced.

There’s another more interesting approach though. We don’t really need the complete inorder traversal. We just need a way to iterate through it from beginning and end, one element at a time. If we had a pointer to a BST node and could ask for its inorder successor or predecessor, we’d be done, right?

Normally, given just a BST node, we would need pointer to the parent node to be able to efficiently compute its inorder successor or predecessor, and it would still have O(h) time complexity, where h is the height of the BST. But if we are doing this for all the nodes in order, we can maintain enough state to amortize this cost to constant work per node.

We can do this with a stack. If we maintain the loop invariant that the stack always has the complete path, from the next element, in inorder traversal, to the root, in reverse order (i.e., with the next element at the top), we can get the next element in constant time, and (if it has a right subtree) push all the elements on the path to its inorder successor, to maintain the loop invariant for the next iteration. Since each node gets added and removed just once from the stack, the overall time complexity is linear in the size of the tree. The space complexity is O(h), which would be logarithmic in the size of the tree if it was balanced.

Code might explain it better than the description. In languages that provide the syntactic sugar of generators or enumerators, the code can be quite pretty. Otherwise, you’ll have to implement the generator pattern on your own: create a class to store state and provide the equivalent of hasNext() and getNext() methods.

Let’s create a simple BST node:

This is how we can implement the inorder traversal from start using generators in python:

We also need a way to iterate through inorder traversal from the end.

Now that we’ve a way to iterate through the inorder traversal from both beginning and end, we can just solve it with the same two pointer approach as for a sorted array.

Notice the similarity with the code for a sorted array now that we’ve hidden away how we are traversing, using a generator.

You can look at the whole code here.