在Python中实现深度优先树迭代器

Question

我正在尝试在Python中为非二叉树实现一个迭代器类。构造迭代器后，可以重复调用其next()函数以深度优先顺序遍历树（例如此顺序），最后当没有节点时返回None。

下面是树的基本Node类：

class Node(object):

    def __init__(self, title, children=None):
        self.title = title
        self.children = children or []
        self.visited = False   

    def __str__(self):
        return self.title

如上所示，我在第一种方法中为节点引入了一个visited属性，因为我没有看到其他的解决方法。有了这个额外的状态，Iterator类就像这样：

class Iterator(object):

    def __init__(self, root):
        self.stack = []
        self.current = root

    def next(self):
        if self.current is None:
            return None

        self.stack.append(self.current)
        self.current.visited = True

        # Root case
        if len(self.stack) == 1:
            return self.current

        while self.stack:
            self.current = self.stack[-1] 
            for child in self.current.children:
                if not child.visited:
                    self.current = child
                    return child

            self.stack.pop()

这些都很好，但我想摆脱对visited属性的需求，而不使用递归或任何其他对Node类的修改。

我需要的所有状态应该在迭代器中处理，但我不知道如何做到这一点。为整个树保留已访问列表是不可扩展的，也是不可能的，因此必须有一个巧妙的方法来使用堆栈。

尤其令我困惑的是这一点——由于next()函数当然会返回，那么我怎么能记住我去过哪里，而又不标记任何东西或使用过量的存储空间呢？直觉上，我想循环遍历子节点，但是当next()函数返回时，这种逻辑就被打破/忘记了！

更新-这里是一个小测试：

tree = Node(
    'A', [
        Node('B', [
            Node('C', [
                Node('D')
                ]),
            Node('E'),
            ]),
        Node('F'),
        Node('G'),
        ])

iter = Iterator(tree)

out = object()
while out:
    out = iter.next()
    print out

- norman

---

保持一个已访问的列表可能不具有可扩展性，但是基于节点对象ID的已访问集合呢？ - michaelb

那仍然有可能包含每个标签。我希望迭代器一次只保留树的子集。 - norman

“小测试”的预期输出是什么？ - Robᵩ

它应该返回 A B C D E F G None。顺便说一下，实现 @mgilson 的生成器解决方案作为 next() 主体会导致无限循环，但这可能只是我对生成器的适应/理解不足。 - norman

我不确定你是如何尝试做到这一点的，但我已经更新了我的答案以展示它通过了你的测试...（请注意，我编写代码时假设你会在for循环中遍历它，而不是某种while循环。我没有看到你的while循环在任何地方起作用，因为你没有捕获StopIteration异常。） - mgilson

Answer 1

如果您真的需要避免递归，那么这个迭代器可以使用：

from collections import deque

def node_depth_first_iter(node):
    stack = deque([node])
    while stack:
        # Pop out the first element in the stack
        node = stack.popleft()
        yield node
        # push children onto the front of the stack.
        # Note that with a deque.extendleft, the first on in is the last
        # one out, so we need to push them in reverse order.
        stack.extendleft(reversed(node.children))

---

话虽如此，我认为你在考虑这个问题时想得太复杂了。一个好用的（递归）生成器也能解决这个问题：

class Node(object):

    def __init__(self, title, children=None):
        self.title = title
        self.children = children or []

    def __str__(self):
        return self.title

    def __iter__(self):
        yield self
        for child in self.children:
            for node in child:
                yield node

这两个都通过了你的测试：

expected = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
# Test recursive generator using Node.__iter__
assert [str(n) for n in tree] == expected

# test non-recursive Iterator
assert [str(n) for n in node_depth_first_iter(tree)] == expected

如果你喜欢，你可以很容易地使Node.__iter__使用非递归形式：

def __iter__(self):
   return node_depth_first_iter(self)

- mgilson

---

将“stack = queue”及其左侧的所有操作转换为队列。为什么不只使用堆栈和弹出/追加操作呢？ - kAldown

Answer 2

但是你现在仍然可能会保存每个标签。我希望迭代器一次只保留树的子集。

但是你已经持有了所有东西。请记住，对象本质上是一个字典，每个属性都有一个条目。在Node的__init__中有self.visited = False意味着无论如何，您都会为每个单独的Node对象存储冗余的“visited”键和False值。至少，集合也有可能不保存每个节点ID。试试这个：

class Iterator(object):
    def __init__(self, root):
        self.visited_ids = set()
        ...

    def next(self):
        ...
        #self.current.visited = True
        self.visited_ids.add(id(self.current))
        ...
                #if not child.visited:
                if id(child) not in self.visited_ids:

在集合中查找ID应该与访问节点属性一样快。这种方法比你的解决方案更浪费的唯一方式是集合对象本身的开销（而不是它的元素），只有当您有多个并发迭代器时才需要考虑这一点（显然您没有，否则节点“visited”属性对您将没有用处）。

- nmclean