迭代器

简介

迭代器对象可以在 for 循环中使用:

In [1]:

  1. x = [2, 4, 6]
  2.  
  3. for n in x:
  4. print n
  1. 2
  2. 4
  3. 6

其好处是不需要对下标进行迭代,但是有些情况下,我们既希望获得下标,也希望获得对应的值,那么可以将迭代器传给 enumerate 函数,这样每次迭代都会返回一组 (index, value) 组成的元组:

In [2]:

  1. x = [2, 4, 6]
  2.  
  3. for i, n in enumerate(x):
  4. print 'pos', i, 'is', n
  1. pos 0 is 2
  2. pos 1 is 4
  3. pos 2 is 6

迭代器对象必须实现 iter 方法:

In [3]:

  1. x = [2, 4, 6]
  2. i = x.__iter__()
  3. print i
  1. <listiterator object at 0x0000000003CAE630>

iter() 返回的对象支持 next 方法,返回迭代器中的下一个元素:

In [4]:

  1. print i.next()
  1. 2

当下一个元素不存在时,会 raise 一个 StopIteration 错误:

In [5]:

  1. print i.next()
  2. print i.next()
  1. 4
  2. 6

In [6]:

  1. i.next()
  1. ---------------------------------------------------------------------------
  2. StopIteration Traceback (most recent call last)
  3. <ipython-input-6-e590fe0d22f8> in <module>()
  4. ----> 1 i.next()
  5.  
  6. StopIteration:

很多标准库函数返回的是迭代器:

In [7]:

  1. r = reversed(x)
  2. print r
  1. <listreverseiterator object at 0x0000000003D615F8>

调用它的 next() 方法:

In [8]:

  1. print r.next()
  2. print r.next()
  3. print r.next()
  1. 6
  2. 4
  3. 2

字典对象的 iterkeys, itervalues, iteritems 方法返回的都是迭代器:

In [9]:

  1. x = {'a':1, 'b':2, 'c':3}
  2. i = x.iteritems()
  3. print i
  1. <dictionary-itemiterator object at 0x0000000003D51B88>

迭代器的 iter 方法返回它本身:

In [10]:

  1. print i.__iter__()
  1. <dictionary-itemiterator object at 0x0000000003D51B88>

In [11]:

  1. print i.next()
  1. ('a', 1)

自定义迭代器

自定义一个 list 的取反迭代器:

In [12]:

  1. class ReverseListIterator(object):
  2.  
  3. def __init__(self, list):
  4. self.list = list
  5. self.index = len(list)
  6.  
  7. def __iter__(self):
  8. return self
  9.  
  10. def next(self):
  11. self.index -= 1
  12. if self.index >= 0:
  13. return self.list[self.index]
  14. else:
  15. raise StopIteration

In [13]:

  1. x = range(10)
  2. for i in ReverseListIterator(x):
  3. print i,
  1. 9 8 7 6 5 4 3 2 1 0

只要我们定义了这三个方法,我们可以返回任意迭代值:

In [14]:

  1. class Collatz(object):
  2.  
  3. def __init__(self, start):
  4. self.value = start
  5.  
  6. def __iter__(self):
  7. return self
  8.  
  9. def next(self):
  10. if self.value == 1:
  11. raise StopIteration
  12. elif self.value % 2 == 0:
  13. self.value = self.value / 2
  14. else:
  15. self.value = 3 * self.value + 1
  16. return self.value

这里我们实现 Collatz 猜想

  • 奇数 n:返回 3n + 1
  • 偶数 n:返回 n / 2 直到 n 为 1 为止:

In [15]:

  1. for x in Collatz(7):
  2. print x,
  1. 22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1

不过迭代器对象存在状态,会出现这样的问题:

In [16]:

  1. i = Collatz(7)
  2. for x, y in zip(i, i):
  3. print x, y
  1. 22 11
  2. 34 17
  3. 52 26
  4. 13 40
  5. 20 10
  6. 5 16
  7. 8 4
  8. 2 1

一个比较好的解决方法是将迭代器和可迭代对象分开处理,这里提供了一个二分树的中序遍历实现:

In [17]:

  1. class BinaryTree(object):
  2. def __init__(self, value, left=None, right=None):
  3. self.value = value
  4. self.left = left
  5. self.right = right
  6.  
  7. def __iter__(self):
  8. return InorderIterator(self)

In [18]:

  1. class InorderIterator(object):
  2.  
  3. def __init__(self, node):
  4. self.node = node
  5. self.stack = []
  6.  
  7. def next(self):
  8. if len(self.stack) > 0 or self.node is not None:
  9. while self.node is not None:
  10. self.stack.append(self.node)
  11. self.node = self.node.left
  12. node = self.stack.pop()
  13. self.node = node.right
  14. return node.value
  15. else:
  16. raise StopIteration()

In [19]:

  1. tree = BinaryTree(
  2. left=BinaryTree(
  3. left=BinaryTree(1),
  4. value=2,
  5. right=BinaryTree(
  6. left=BinaryTree(3),
  7. value=4,
  8. right=BinaryTree(5)
  9. ),
  10. ),
  11. value=6,
  12. right=BinaryTree(
  13. value=7,
  14. right=BinaryTree(8)
  15. )
  16. )

In [20]:

  1. for value in tree:
  2. print value,
  1. 1 2 3 4 5 6 7 8

不会出现之前的问题:

In [21]:

  1. for x,y in zip(tree, tree):
  2. print x, y
  1. 1 1
  2. 2 2
  3. 3 3
  4. 4 4
  5. 5 5
  6. 6 6
  7. 7 7
  8. 8 8

原文: https://nbviewer.jupyter.org/github/lijin-THU/notes-python/blob/master/05-advanced-python/05.09-iterators.ipynb