【Python Cookbook】迭代器与生成器（三）

9.排列组合的迭代

你想迭代遍历一个集合中元素的所有可能的排列或组合。

itertools 模块提供了三个函数来解决这类问题。其中一个是 itertools.permutations()，它接受一个集合并产生一个元组序列，每个元组由集合中所有元素的一个可能排列组成。也就是说通过打乱集合中元素排列顺序生成一个元组，比如：

>>> items = ['a', 'b', 'c']
>>> from itertools import permutations
>>> for p in permutations(items):
...     print(p)
...
('a', 'b', 'c')
('a', 'c', 'b')
('b', 'a', 'c')
('b', 'c', 'a')
('c', 'a', 'b')
('c', 'b', 'a')
>>>

如果你想得到指定长度的所有排列，你可以传递一个可选的长度参数。就像这样：

>>> for p in permutations(items, 2):
...     print(p)
...
('a', 'b')
('a', 'c')
('b', 'a')
('b', 'c')
('c', 'a')
('c', 'b')
>>>

使用 itertools.combinations() 可得到输入集合中元素的所有的组合。比如：

>>> from itertools import combinations
>>> for c in combinations(items, 3):
...     print(c)
...
('a', 'b', 'c')

>>> for c in combinations(items, 2):
...     print(c)
...
('a', 'b')
('a', 'c')
('b', 'c')

>>> for c in combinations(items, 1):
...     print(c)
...
('a',)
('b',)
('c',)
>>>

对于 combinations() 来讲，元素的顺序已经不重要了。 也就是说，组合 ('a', 'b') 跟 ('b', 'a') 其实是一样的（最终只会输出其中一个）。

在计算组合的时候，一旦元素被选取就会从候选中剔除掉（比如如果元素 a 已经被选取了，那么接下来就不会再考虑它了）。而函数 itertools.combinations_with_replacement() 允许同一个元素被选择多次，比如：

>>> for c in combinations_with_replacement(items, 3):
...     print(c)
...
('a', 'a', 'a')
('a', 'a', 'b')
('a', 'a', 'c')
('a', 'b', 'b')
('a', 'b', 'c')
('a', 'c', 'c')
('b', 'b', 'b')
('b', 'b', 'c')
('b', 'c', 'c')
('c', 'c', 'c')
>>>

这一小节我们向你展示的仅仅是 itertools 模块的一部分功能。尽管你也可以自己手动实现排列组合算法，但是这样做得要花点脑力。当我们碰到看上去有些复杂的迭代问题时，最好可以先去看看 itertools 模块。如果这个问题很普遍，那么很有可能会在里面找到解决方案！

10.序列上索引值迭代

你想在迭代一个序列的同时跟踪正在被处理的元素索引。

内置的 enumerate() 函数可以很好的解决这个问题：

>>> my_list = ['a', 'b', 'c']
>>> for idx, val in enumerate(my_list):
...     print(idx, val)
...
0 a
1 b
2 c

为了按传统行号输出（行号从 1 开始），你可以传递一个开始参数：

>>> my_list = ['a', 'b', 'c']
>>> for idx, val in enumerate(my_list, 1):
...     print(idx, val)
...
1 a
2 b
3 c

这种情况在你遍历文件时想在错误消息中使用行号定位时候非常有用：

def parse_data(filename):
    with open(filename, 'rt') as f:
        for lineno, line in enumerate(f, 1):
            fields = line.split()
            try:
                count = int(fields[1])
                ...
            except ValueError as e:
                print('Line {}: Parse error: {}'.format(lineno, e))

enumerate() 对于跟踪某些值在列表中出现的位置是很有用的。所以，如果你想将一个文件中出现的单词映射到它出现的行号上去，可以很容易的利用 enumerate() 来完成：

word_summary = defaultdict(list)

with open('myfile.txt', 'r') as f:
    lines = f.readlines()

for idx, line in enumerate(lines):
    # Create a list of words in current line
    words = [w.strip().lower() for w in line.split()]
    for word in words:
        word_summary[word].append(idx)

如果你处理完文件后打印 word_summary，会发现它是一个字典（准确来讲是一个 defaultdict）， 对于每个单词有一个 key ，每个 key 对应的值是一个由这个单词出现的行号组成的列表。如果某个单词在一行中出现过两次，那么这个行号也会出现两次，同时也可以作为文本的一个简单统计。

当你想额外定义一个计数变量的时候，使用 enumerate() 函数会更加简单。你可能会像下面这样写代码：

lineno = 1
for line in f:
    # Process line
    ...
    lineno += 1

但是如果使用 enumerate() 函数来代替就显得更加优雅了：

for lineno, line in enumerate(f):
    # Process line
    ...

enumerate() 函数返回的是一个 enumerate 对象实例，它是一个迭代器，返回连续的包含一个计数和一个值的元组，元组中的值通过在传入序列上调用 next() 返回。

还有一点可能并不很重要，但是也值得注意，有时候当你在一个已经解压后的元组序列上使用 enumerate() 函数时很容易调入陷阱。你得像下面正确的方式这样写：

data = [ (1, 2), (3, 4), (5, 6), (7, 8) ]

# Correct!
for n, (x, y) in enumerate(data):
    ...
# Error!
for n, x, y in enumerate(data):
    ...

11.同时迭代多个序列

你想同时迭代多个序列，每次分别从一个序列中取一个元素。

为了同时迭代多个序列，使用 zip() 函数。比如：

>>> xpts = [1, 5, 4, 2, 10, 7]
>>> ypts = [101, 78, 37, 15, 62, 99]
>>> for x, y in zip(xpts, ypts):
...     print(x,y)
...
1 101
5 78
4 37
2 15
10 62
7 99
>>>

zip(a, b) 会生成一个可返回元组 (x, y) 的迭代器，其中 x 来自 a，y 来自 b。一旦其中某个序列到底结尾，迭代宣告结束。因此迭代长度跟参数中最短序列长度一致。

>>> a = [1, 2, 3]
>>> b = ['w', 'x', 'y', 'z']
>>> for i in zip(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
>>>

如果这个不是你想要的效果，那么还可以使用 itertools.zip_longest() 函数来代替。比如：

>>> from itertools import zip_longest
>>> for i in zip_longest(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(None, 'z')

>>> for i in zip_longest(a, b, fillvalue=0):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(0, 'z')
>>>

当你想成对处理数据的时候 zip() 函数是很有用的。比如，假设你头列表和一个值列表，就像下面这样：

headers = ['name', 'shares', 'price']
values = ['ACME', 100, 490.1]

使用 zip() 可以让你将它们打包并生成一个字典：

s = dict(zip(headers,values))

或者你也可以像下面这样产生输出：

for name, val in zip(headers, values):
    print(name, '=', val)

虽然不常见，但是 zip() 可以接受多于两个的序列的参数。这时候所生成的结果元组中元素个数跟输入序列个数一样。比如;

>>> a = [1, 2, 3]
>>> b = [10, 11, 12]
>>> c = ['x','y','z']
>>> for i in zip(a, b, c):
...     print(i)
...
(1, 10, 'x')
(2, 11, 'y')
(3, 12, 'z')
>>>

最后强调一点就是， zip() 会创建一个迭代器来作为结果返回。如果你需要将结对的值存储在列表中，要使用 list() 函数。比如：

>>> zip(a, b)
<zip object at 0x1007001b8>
>>> list(zip(a, b))
[(1, 10), (2, 11), (3, 12)]
>>>

12.不同集合上元素的迭代

你想在多个对象执行相同的操作，但是这些对象在不同的容器中，你希望代码在不失可读性的情况下避免写重复的循环。

itertools.chain() 方法可以用来简化这个任务。它接受一个可迭代对象列表作为输入，并返回一个迭代器，有效的屏蔽掉在多个容器中迭代细节。为了演示清楚，考虑下面这个例子：

>>> from itertools import chain
>>> a = [1, 2, 3, 4]
>>> b = ['x', 'y', 'z']
>>> for x in chain(a, b):
... print(x)
...
1
2
3
4
x
y
z
>>>

使用 chain() 的一个常见场景是当你想对不同的集合中所有元素执行某些操作的时候。比如：

# Various working sets of items
active_items = set()
inactive_items = set()

# Iterate over all items
for item in chain(active_items, inactive_items):
    # Process item

这种解决方案要比像下面这样使用两个单独的循环更加优雅，

for item in active_items:
    # Process item
    ...

for item in inactive_items:
    # Process item
    ...

itertools.chain() 接受一个或多个可迭代对象作为输入参数。然后创建一个迭代器，依次连续的返回每个可迭代对象中的元素。这种方式要比先将序列合并再迭代要高效的多。比如：

# Inefficent
for x in a + b:
    ...

# Better
for x in chain(a, b):
    ...

第一种方案中， a + b 操作会创建一个全新的序列并要求 a 和 b 的类型一致。chian() 不会有这一步，所以如果输入序列非常大的时候会很省内存。并且当可迭代对象类型不一样的时候 chain() 同样可以很好的工作。