Lazy Graph Traversal

When we need to traverse a graph, we tend to choose depth-first-search because of its simplicity. However, sometimes it can be desirable to pause and resume the traversal whenever we want. In this post, we will explore how to do graph traversal and control when the next node is visited.

Prerequisite: A basic understanding of Depth First Search is enough.

Setup

To make an example, we will create a tree using this node class and do an in-order traversal over it.

1class Node:
2    def __init__(self, val=0, left=None, right=None):
3        self.val = val
4        self.left = left
5        self.right = right

Using an Explicit Stack

The problem with the usual depth-first search is that we cannot pause it. The recursive function calls itself, so we cannot easily tell it to pause before moving on to the next node. We can get around this issue by using a stack that serves the same purpose as the implicit function call stack.

1class LazyExplicitStackBasedExplorer:
2    def __init__(self, root):
3        self.stack = []
4        self.current_node = root
5        self.next_value = None
6
7    def get_next_value(self):
8        if self.current_node is not None or len(self.stack) > 0:
9            self.visit_next_node()
10            return self.next_value
11        else:
12            return None
13
14    def visit_next_node(self):
15        while self.current_node is not None:
16            self.stack.append(self.current_node)
17            self.current_node = self.current_node.left
18
19        self.current_node = self.stack.pop()
20        self.next_value = self.current_node.val
21        self.current_node = self.current_node.right

Using Generators

Using an explicit stack is fine, but it requires us to maintain the stack ourselves. Can we make it as simple as a recursive DFS solution? Luckily some languages support generator functions. Python has it. JavaScript also has it. These are like functions, but we can pause their execution to return some value and then resume the execution when we want the next value. In our case, we can make a recursive generator function that visits a node, returns it, and pauses the execution until we tell it to get the next node.

1class LazyGeneratorBasedExplorer:
2    def __init__(self, root):
3        self.iterator = self.visit_next_node(root)
4
5    def get_next_value(self):
6        try:
7            return next(self.iterator)
8        except StopIteration:
9            return None
10
11    def visit_next_node(self, root):
12        if root.left is not None:
13            for value in self.visit_next_node(root.left):
14                yield value
15
16        yield root.val
17
18        if root.right is not None:
19            for value in self.visit_next_node(root.right):
20                yield value

Use case

The Leet Code problem #653. Two Sum IV - Input is a BST asks us to find two nodes that add up to a target value inside a binary search tree. This problem sounds a lot like Two Sum II - Input Array Is Sorted, but the input is a tree instead of an array.

We could do a simple DFS in-order traversal of the tree, create a sorted array of the nodes and solve it using two-pointers. It works but requires O(n) auxiliary space for a tree with n nodes.

Using our lazy graph traversal, we can solve the problem without traversing the entire tree and without using extra data structures (Not counting the recursion stack). My solution: We start two lazy in-order traversals of the tree, one from the leftmost leaf node and another from the rightmost leaf node. Now, just like the two-pointers approach, we compare our current left and right values (nodes in this case) and advance one side at a time

Online playground: Example in ideone.com