tonysdg tonysdg - 7 months ago 14
Python Question

Can you define multiple different iterators for a Python class?

I'm writing a very simple Tree class:

class Tree:
def __init__(self, value_ = None, children_ = None):
self.value = value_
self.children = children_


I'd like to be able to perform both DFS and BFS traversal with a simple loop, i.e.:

t = Tree()
# ...fill tree...

for node in t:
print(node.value)


In C++, for example, you can have multiple types of iterators - so I could define both a DFS and a BFS iterator and use one or the other depending on what type of traversal I wanted to do. Is this possible to do in Python?

Answer

You can have multiple methods returning iterators and have the 'default' one as __iter__. Below is a simple binary tree where 'default' iterator does DFS and which additionally supports BFS with separate method:

from collections import deque

class Tree(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def __iter__(self):
        if self.left:
            for x in self.left:
                yield x

        yield self.value

        if self.right:
            for x in self.right:
                yield x

    def bfs(self):
        q = deque([self])
        while q:
            x = q.popleft()
            if x:
                yield x.value
                q.extend([x.left, x.right])

Short example of usage:

root = Tree(2)
root.left = Tree(1)
root.right = Tree(4)
root.right.left = Tree(3)
root.right.right = Tree(5)

print list(root) # [1, 2, 3, 4, 5]
print list(root.bfs()) # [2, 1, 4, 3, 5]