python的迭代器和生成器

1、迭代器

简介

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

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

输出：

2
4
6
1
2
3

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

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

输出：

pos 0 is 2
pos 1 is 4
pos 2 is 6
1
2
3

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

x = [2, 4, 6]
i = x.__iter__()
print i
1
2
3

输出：

<listiterator object at 0x0000000003CAE630>
1

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

print i.next()
1

输出：

2
1

当下一个元素不存在时，会 raise 一个 StopIteration 错误：
如：

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

输出：

4
6
1
2

自定义迭代器

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

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

x = range(10)
for i in ReverseListIterator(x):
    print i,
1
2
3

输出：

9 8 7 6 5 4 3 2 1 0
1

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

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

这里我们实现 Collatz 猜想：
奇数 n：返回 3n + 1
偶数 n：返回 n / 2
直到 n 为 1 为止：

for x in Collatz(7):
    print x,
1
2

输出：

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1
1

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

i = Collatz(7)
for x, y in zip(i, i):
    print x, y
1
2
3

输出：

22 11
34 17
52 26
13 40
20 10
5 16
8 4
2 1
1
2
3
4
5
6
7
8

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

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

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

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

for value in tree:
    print value,
1
2

输出：

1 2 3 4 5 6 7 8
1

不会出现之前的问题：

for x,y in zip(tree, tree):
    print x, y
1
2

1 1
2 2
3 3
4 4
5 5
6 6
7 7
8 8
1
2
3
4
5
6
7
8

生成器

简介

while 循环通常有这样的形式：

<do setup>
result = []
while True:
    <generate value>
    result.append(value)
    if <done>:
        break
1
2
3
4
5
6
7

使用迭代器实现这样的循环：

class GenericIterator(object):
    def __init__(self, ...):
        <do setup>
        # 需要额外储存状态
        <store state>
    def next(self): 
        <load state>
        <generate value>
        if <done>:
            raise StopIteration()
        <store state>
        return value
1
2
3
4
5
6
7
8
9
10
11
12

更简单的，可以使用生成器：

def generator(...):
    <do setup>
    while True:
        <generate value>
        # yield 说明这个函数可以返回多个值！
        yield value
        if <done>:
            break
1
2
3
4
5
6
7
8

生成器使用 yield 关键字将值输出，而迭代器则通过 next 的 return 将值返回；与迭代器不同的是，生成器会自动记录当前的状态，而迭代器则需要进行额外的操作来记录当前的状态。

对于之前的 collatz 猜想，简单循环的实现如下：

def collatz(n):
    sequence = []
    while n != 1:
        if n % 2 == 0:
            n /= 2
        else:
            n = 3*n + 1
        sequence.append(n)
    return sequence
 
for x in collatz(7):
    print x,
1
2
3
4
5
6
7
8
9
10
11
12

输出：

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1
1

迭代器的版本如下：

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

输出：

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1
1

生成器的版本如下：

def collatz(n):
    while n != 1:
        if n % 2 == 0:
            n /= 2
        else:
            n = 3*n + 1
        yield n
 
for x in collatz(7):
    print x,
1
2
3
4
5
6
7
8
9
10

输出：

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1
1

事实上，生成器也是一种迭代器：

x = collatz(7)
print x
1
2

输出;

<generator object collatz at 0x0000000003B63750>
1

它支持 next 方法，返回下一个 yield 的值：

print x.next()
print x.next()
1
2

22
11
1
2

iter 方法返回的是它本身：

print x.__iter__()
1

<generator object collatz at 0x0000000003B63750>
1

之前的二叉树迭代器可以改写为更简单的生成器模式来进行中序遍历：

class BinaryTree(object):
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right
 
    def __iter__(self):
        # 将迭代器设为生成器方法
        return self.inorder()
 
    def inorder(self):
        # traverse the left branch
        if self.left is not None:
            for value in self.left:
                yield value
 
        # yield node's value
        yield self.value
 
        # traverse the right branch
        if self.right is not None:
            for value in self.right:
                yield value
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

非递归的实现：

def inorder(self):
    node = self
    stack = []
    while len(stack) > 0 or node is not None:
        while node is not None:
            stack.append(node)
            node = node.left
        node = stack.pop()
        yield node.value
        node = node.right
1
2
3
4
5
6
7
8
9
10

tree = BinaryTree(
    left=BinaryTree(
        left=BinaryTree(1),
        value=2,
        right=BinaryTree(
            left=BinaryTree(3),
            value=4,
            right=BinaryTree(5)
        ),
    ),
    value=6,
    right=BinaryTree(
        value=7,
        right=BinaryTree(8)
    )
)
for value in tree:
    print value,
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

1 2 3 4 5 6 7 8
1