graph_drawing/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()