Intuitive Iterators For Binary Trees

As a follow-up to the previous article on intuitive iterative traversals for binary trees, we discuss the application of the intuitive stack based method to the task of creating iterators for binary trees. The iterator is an abstraction for making incremental traversals over data containers or data streams. At its core, the iterator pattern is implemented by two functions:

• next - a function that retrieves the next element in the data container or stream and advances the underlying cursor
• hasNext - a function that indicates whether or not the underlying data container or stream has been exhausted.

As a design pattern, the iterator can be used for accessing data streams with an arbitrary, and sometimes infinite, number of elements. For data structures, an iterator is useful for performing computations based on advancing cursors in certain defined orders. An elementary application of iterators to binary trees is for traversing a given tree in preorder, inorder, or postorder fashion one element at a time. For example, a specific use case might be for an iterator over binary search trees that provide access to comparable elements in increasing order by means of an inorder iterator.

Implementing an iterator for binary trees by modeling the access pattern for a traversal is not a very intuitive task. However, using the technique covered in the previous article of a stack and an algorithm that employs bookkeeping of discovered and undiscovered nodes, an implementation follows quite naturally.

Let’s walk through implementation of an inorder iterator. To implement this iterator, its internal state can be initialized in the same way as the beginning of a conventional, all-at-once iterative traversal by pushing the root node onto an empty stack.

def __init__(self, root):
    self.stack = []
    self.discovered = {}
    if root:
        self.stack.append(root)

The stack will serve as the primary internal state of the iterator. When the stack is empty, that indicates the end of the traversal and the depletion of elements.

def hasNext(self):
    return len(self.stack) > 0

Traversing a tree while advancing the implicit cursor requires a bit of thought. Just like in the conventional, all-at-once traversal, the top of the stack always contains the next element to be processed. However, it may not be the next element to be visited. For inorder traversals, we previously determined that a visit should happen only if the node has already been discovered. Thus, we will continue to process nodes until we come across a discovered node, in which case its value can be returned. Any undiscovered nodes should be processed by marking it as discovered and pushing its children onto the stack in the appropriate order.

def next(self):
    node = self.stack.pop()
    while node not in self.discovered:
        self.discovered[node] = True
        if node.right:
            self.stack.append(node.right)
        self.stack.append(node)
        if node.left:
            self.stack.append(node.left)
        node = self.stack.pop()
    return node.val

As a sanity check, we can walk through a base case using the logic above. For a single node binary tree, the root node will initially be undiscovered. The first invocation of next will pop the root node off the stack, mark it as discovered, then pop it off the stack again. Since the node is now marked as discovered, its value can be returned. The resulting stack is empty and an invocation of hasNext will return false, indicating that the iterator has been exhausted.

Putting it all together, we have an implementation of an inorder iterator for binary trees that follows quite naturally from the intuitive stack based method of traversing binary trees.

class BinaryTreeInorderIterator:
    def __init__(self, root):
        self.stack = []
        self.discovered = {}
        if root:
            self.stack.append(root)

    def hasNext(self):
        return len(self.stack) > 0

    def next(self):
        node = self.stack.pop()
        while node not in self.discovered:
            self.discovered[node] = True
            if node.right:
                self.stack.append(node.right)
            self.stack.append(node)
            if node.left:
                self.stack.append(node.left)
            node = self.stack.pop()
        return node.val

Like previously discussed, an implementation of a preorder or postorder iterator should easily be derived by modifying the order in which a node and its child nodes are pushed onto the stack.