Search code examples
python-3.xgeneratorbreadth-first-searchyield-from

Breadth First Tree Traversal using Generators in Python


I am studying how to use Generators in Python in David Beazly's excellent Python Cookbook text. The following code recipe defines Depth First Tree Traversal using generators very elegantly:

# example.py
#
# Example of depth-first search using a generator

class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children)

    def depth_first(self):
        yield self
        for c in self:
            yield from c.depth_first()

# Example
if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    child1.add_child(Node(3))
    child1.add_child(Node(4))
    child2.add_child(Node(5))

    for ch in root.depth_first():
        print(ch)
    # Outputs: Node(0), Node(1), Node(3), Node(4), Node(2), Node(5)

I am trying to come up with an equally elegant method

def breadth_first(self):
    pass

I am deliberately not posting the crazy stuff that I have been trying since everything that I have tried requires maintaining 'state' within it. I don't want to use the traditional queue based solutions. The whole point of this academic exercise is to learn how generators behave in depth. Therefore, I want to create a parallel 'breadth_first' method using generators for the tree above.

Any pointers/solutions are welcome.


Solution

  • You can't use recursion (stack) for bfs without some serious hacks, but a queue would work:

    def breadth_first(self):
        q = [self]
        while q:
            n = q.pop(0)
            yield n
            for c in n._children:
                q.append(c)