你能为Python类定义多个不同的迭代器吗？

Question

我正在写一个非常简单的Tree类:

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

我希望能够通过一个简单的循环执行DFS和BFS遍历,即:

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

for node in t:
    print(node.value)

例如,在C++中,您可以拥有多种类型的迭代器 - 因此我可以定义DFS和BFS迭代器,并根据我想要执行的遍历类型使用其中一种.这可以用Python做吗？

Answer 1

您可以使用多个方法返回迭代器,并将"默认"方法作为__iter__.下面是一个简单的二叉树,其中'default'迭代器执行DFS,并且还使用单独的方法支持BFS:

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])

用法的简短示例:

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]