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