Added iterative depth-first traversal routine.

traverse.py

+"""
+Tree traversal routines.
+
+Created by Sander Mathijs Van Veen <smvv@kompiler.org>, 2012.
+"""
+from node import Leaf
+
+
+def traverse_depth_first(root):
+    """
+    Traverse the given root of a tree in depth-first order. This is an
+    iterative instead of the ordinary recursive implementation. It uses
+    iterators to reduce memory usage and, since it is iterative, the depth of
+    the given tree does not hit Python's recursion limit.
+
+    This implementation uses two stacks. One contains the path from the root
+    node to the current visited node / leaf using child node indices. The other
+    one contains pointers to the nodes which are not yet visited (aka still on
+    the "to do" list). Thus, traversal using this implementation has a worst
+    case space complexity of O(2 * log N), where N is the number of nodes in
+    the tree.
+
+    >>> from node import Node as N, Leaf as L
+    >>> root = N('+', N('/', L(1), L(2)), N('*', L(3), L(4)))
+    >>> print map(lambda n: n.title(), traverse_depth_first(root))
+    ['1', '2', '/', '3', '4', '*', '+']
+    >>> root = N('+', N('/', L(1), N('-', L(2), N('^', L(3), L(4)))), \
+                      N('*', L(5), L(6)))
+    >>> print map(lambda n: n.title(), traverse_depth_first(root))
+    ['1', '2', '3', '4', '^', '-', '/', '5', '6', '*', '+']
+    """
+    if not root:
+        raise StopIteration()
+
+    # List of nodes which are not yet visited (the "to do" list).
+    stack = []
+
+    # List of child node indices which indicates the path to the current
+    # visited node or leaf.
+    path = []
+
+    node = root
+
+    while True:
+        while not isinstance(node, Leaf):
+            # Traverse left and save the path
+            stack.append(node)
+            path.append(0)
+            node = node[0]
+
+        yield node
+
+        if not stack:
+            break
+
+        while stack:
+            # Set pointer to parent node and increase child node position.
+            node = stack[-1]
+            path[-1] += 1
+            pos = path[-1]
+
+            if pos >= len(node):
+                # Visit the parent node, since all children are visited.
+                yield node
+
+                # Remove the node and move up in the tree.
+                stack.pop()
+                path.pop()
+
+                if not stack:
+                    raise StopIteration()
+            else:
+                # This child node is not visited.
+                node = node[pos]
+                break
+
+    raise StopIteration()
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()