使用列表遍历树的Python迭代器

Question

我对Python中的迭代器还不是很了解。我有一个包含子节点的对象，我想遍历这个结构。我希望能像使用printall函数那样的效果，但用迭代器来实现。

    class t:
            def __init__(self, i):
                    self.l = []
                    self.a = 0
                    for ii in range(i):
                            self.a = ii
                            self.l.append(t(i-1))

            def __iter__(self):
                    return self

            def next(self):
                    for i in self.l:
                            yield i.__iter__()
                    yield self

            def printall(self):
                    for i in self.l:
                            i.printall()
                    print self.a

希望这些信息足够了，谢谢。

补充：

我只是想能够遍历树的所有叶子节点，并对每个对象做点什么，也就是说，当我有一个实例的时候

    bla = t(3) 

我希望能够用

    for x in bla:
            print x.a

比如说，我想对每个x做点什么，我只需要访问每个子节点一次。

Answer 1

我的第一个建议是把你的类名改得更清楚一些，遵循一下PEP-8的规范。像t这样的类名有点难以管理：

class Tree:
    def __init__(self, i):
        self.l = []
        self.a = 0
        for ii in range(i):
            self.a = ii
            self.l.append(Tree(i-1))

接下来，你需要把__iter__()这个方法改成返回self里的下一个元素，而不是返回self本身——这不是开玩笑哦 :) __iter__()方法应该返回一个指向原始对象的迭代器，而不是对象本身：

def __iter__(self):
    return next(self)

现在进入难点：next()方法。我总觉得写递归迭代器很难，但其实也不是不可能：对于每个子节点，遍历它并返回遍历到的值：

def next(self):
    for i in self.l:
        for ii in i:
            yield ii
    yield self

因为这个方法是递归的，所以它会处理所有的子孙节点。当在一个叶子节点（没有子节点的节点）调用next()时，它只会返回这个节点本身。另一方面，当在一个有子节点的节点上调用时，它会对每个子节点调用自己，并返回子节点的值。这意味着它会被子节点的子节点调用，依此类推，直到叶子节点。等到一个节点的所有子孙节点都被调用过——也就是所有子孙节点的值都被返回后，它应该返回它自己的值，所以你需要返回原始节点本身。

现在你的printall()函数应该能正常工作了：

if __name__ == "__main__":
t = Tree(6)
t.printall()

最后的一些注意事项：

• 总是让你的类继承object：

  class Tree(object)::

• 我敢打赌你想写一个像下面这样的__init__()方法：

  def __init__(self, i):
      self.l = []
      self.a = i
      for ii in range(i):
          self.l.append(Tree(i-1))
  

• wberry的解决方案更好，因为它更简洁，也可能更高效。不过，我觉得提问者正在学习树、递归等，所以我认为一个更具体的解决方案会更有教育意义 :)

Answer 2

听起来你想让这个迭代器像树的遍历一样工作。可以研究一下 itertools 模块，这样你会有很多收获。

from itertools import chain, imap

class t:
  def __init__(self, value):
    self.value = value
    self.children = []
  def __iter__(self):
    "implement the iterator protocol"
    for v in chain(*imap(iter, self.children)):
      yield v
    yield self.value

root = t(0)
root.children.append(t(1))
root.children.append(t(2))
root.children[1].children.append(t(3))
print list(iter(root))   # -> [1, 3, 2, 0]
print list(iter(root.children[1]))  # -> [3, 2]

编辑：下面是最初被接受的实现方式。它有性能问题；我本来想删掉它，但觉得删除一个被接受的答案不太合适。这个实现会完全遍历整个结构，创建 O(N*log[M](N)) 个生成器对象（对于一个平衡树，分支因子为 M，总共有 N 个元素），在返回任何值之前。但它确实用一个简单的表达式产生了想要的结果。

（上面的实现是按需访问树的某些区域，同时内存中只会有 O(M+log[M](N)) 个生成器对象。在这两种实现中，预计只有 O(log[M](N)) 层的嵌套生成器。）

from itertools import chain

def isingle(item):
  "iterator that yields only a single value then stops, for chaining"
  yield item

class t:
  # copy __init__ from above
  def __iter__(self):
    "implement the iterator protocol"
    return chain(*(map(iter, self.children) + [isingle(self.value)]))