迭代器模块 itertools

Posted 2021-01-25 famcoo

无限迭代器

itertools 包自带了三个可以无限迭代的迭代器。这意味着，当你使用他们时，你要知道你需要的到底是最终会停止的迭代器，还是需要无限地迭代下去。

这些无限迭代器在生成数字或者在长度未知的可迭代对象（iterables）中循环时相当有用。下面我们开始认识这些有趣的可迭代对象！

count(初值=0, 步长=1)

count 迭代器会返回从传入的起始参数开始的均匀间隔的数值。count 也可以接收指定的步长参数。我们来看一个简单的例子：

>>> from itertools import count
>>> for i in count(10):
...     if i > 20: 
...         break
...     else:
...         print(i)
... 
10
11
12
13
14
15
16
17
18
19
20

这里我们先从 itertools 导入 count，然后创建一个 for 循环。循环中加入了条件检查，当迭代器值大于 20 时跳出循环，否则将打印出迭代器的当前值。你应该注意到了，输出结果从 10 开始，与我们传入 count 的起始值是一致的。

另一种控制无限迭代器输出的方式是使用 itertools 的子模块 islice。使用方法如下：

>>> from itertools import islice
>>> for i in islice(count(10), 5):
...     print(i)
... 
10
11
12
13
14

这里，我们先导入了 islice，然后遍历 count，从 10 开始，输出 5 个元素后结束。你大概猜到了，islice 的第二个参数控制何时停止迭代。但其含义并不是”达到数字 5 时停止“，而是”当迭代了 5 次之后停止“。

cycle(可迭代对象)

itertools 中的 cycle 迭代器允许你创建一个能在一组值间无限循环的迭代器。下面我们传入一个 3 个字母的字符串，看看将会发生什么：

>>> from itertools import cycle
>>> count = 0
>>> for item in cycle(‘XYZ‘):
...     if count > 7:
...         break
...     print(item)
...     count += 1
... 
X
Y
Z
X
Y
Z
X
Y

这里我们创建了一个 for 循环，使其在三个字母 XYZ 间无限循环。当然，我们并不真地想要永远循环下去，所以我们添加了一个简单的计数器来跳出循环。

你也可以用 Python 内建的 next 函数对 itertools 创建的迭代器进行循环：

>>> polys = [‘triangle‘, ‘square‘, ‘pentagon‘, ‘rectangle‘]
 >>> iterator = cycle(polys)
 >>> next(iterator)
 ‘triangle‘
 >>> next(iterator)
 ‘square‘
 >>> next(iterator)
 ‘pentagon‘
 >>> next(iterator)
 ‘rectangle‘
 >>> next(iterator)
 ‘triangle‘
 >>> next(iterator)
 ‘square‘

上面的代码中，我们创建一个简单的多边形组成的列表，然后传入 cycle。我们用一个变量存储新建的迭代器，然后将这个变量传入 next 函数。每一次调用 next 都将返回迭代器中的下一个值。由于迭代器是无限的，我们可以一直调用 next 而永远不会到尽头。

repeat(对象[, 次数])

repeat 迭代器会一遍遍地返回传入的对象直至永远，除非你设定了 times 参数。这与 cycle 非常相像，除了一点，repeat 不会在多个值之间循环。我们看一个简单的例子：

>>> from itertools import repeat
>>> repeat(5, 5)
repeat(5, 5)
>>> iterator = repeat(5, 5)
>>> next(iterator)
5
>>> next(iterator)
5
>>> next(iterator)
5
>>> next(iterator)
5
>>> next(iterator)
5
>>> next(iterator)
Traceback (most recent call last):
  Python Shell, prompt 21, line 1
builtins.StopIteration:

这里，我们先导入 repeat，然后设定其重复五次数字 5。接着，我们连续六次调用 next 函数，观察它是否工作正常。当运行这段代码，将引发 StopIteration，因为最后一次调用 next 函数，迭代器中的值已经全部返回了。

---

可终止的迭代器

多数用 itertools 创建的迭代器并不是无限的。这一节，我们将了解 itertools 中的有限迭代器。为了对于输出的可读的结果，我们将使用 Python 内建列表存储。 如果不使用列表，你将只能打印出 itertools 对象。

accumulate(可迭代对象[, 函数])

accumulate 迭代器将返回累计求和结果，或者传入两个参数的话，由传入的函数累积计算的结果。默认设定为相加，我们赶快试一试吧：

>> from itertools import accumulate
>>> list(accumulate(range(10)))
[0, 1, 3, 6, 10, 15, 21, 28, 36, 45]

这里，我们 导入了 accumulate，然后传入 10 个数字，0-9。迭代器将传入数字依次累加，所以第一个是 0 ，第二个是 0+1， 第三个是 1+2，如此下去。现在我们导入 operator 模块，然后添加进去：

>>> import operator
>>> list(accumulate(range(1, 5), operator.mul))
[1, 2, 6, 24]

这里我们传入了数字 1-4 到 accumulate 迭代器中。我们还传入了一个函数：operator.mul，这个函数将接收的参数相乘。所以每一次迭代，迭代器将以乘法代替除法（1×1=1, 1×2=2, 2×3=6, 以此类推）。

accumulate 的文档中给出了其他一些有趣的例子，例如贷款分期偿还，混沌递推关系等。这绝对值得你花时间去看一看。

chain(*可迭代对象)

chain 迭代器能够将多个可迭代对象合并成一个更长的可迭代对象。实际上，我参与的一个项目中最近就需要这一功能。我有一个列表，里面已经包含一些元素，接着想把另外两个列表添加到最初那个列表中。注意，我们想添加的是两个列表的元素。最初，我是这样做的：

>>> my_list = [‘foo‘, ‘bar‘]
>>> numbers = list(range(5))
>>> cmd = [‘ls‘, ‘/some/dir‘]
>>> my_list.extend(cmd, numbers)
>>> my_list
[‘foo‘, ‘bar‘, [‘ls‘, ‘/some/dir‘], [0, 1, 2, 3, 4]]

这并不是我想要的。itertools 模块提供一个优雅得多的方法用chain 来合并这些列表：

>>> from itertools import chain
>>> my_list = list(chain([‘foo‘, ‘bar‘], cmd, numbers))
>>> my_list
[‘foo‘, ‘bar‘, ‘ls‘, ‘/some/dir‘, 0, 1, 2, 3, 4]

许多聪明的读者可能想到了，实际上不使用 itertools，也有其他方法能够实现这一要求。你可以这样做：

>>> my_list = [‘foo‘, ‘bar‘]
>>> my_list += cmd + numbers
>>> my_list
[‘foo‘, ‘bar‘, ‘ls‘, ‘/some/dir‘, 0, 1, 2, 3, 4]

这些方法当然都是可行的。在我知道 chain 之前，我可能会这样做，但我个人认为这个例子中， chain 更为优雅，也更容易理解。

chain.from_iterable(可迭代对象)

你也可以用 chain 的一个方法，叫做 from_iterable。这个方法与直接用 chain 有些细微的差别。不同于直接传入一系列可迭代对象，你必须传入一个嵌套的列表。我们这就来看一看：

>>> from itertools import chain
>>> numbers = list(range(5))
>>> cmd = [‘ls‘, ‘/some/dir‘]
>>> chain.from_iterable(cmd, numbers)
Traceback (most recent call last):
  Python Shell, prompt 66, line 1
builtins.TypeError: from_iterable() takes exactly one argument (2 given)
>>> list(chain.from_iterable([cmd, numbers]))
[‘ls‘, ‘/some/dir‘, 0, 1, 2, 3, 4]

这里我们跟之前一样先导入 chain。我们尝试传入两个列表，但却出现了 TypeError！为了修正这个错误，我们将 cmd 和 numbers 放入一个列表中，将这个嵌套的列表传入 from_iterable。虽然有点细微差别，但也是很方便使用的。

compress(数据, 选择器)

子模块 compress 在用一个可迭代对象过滤另一个可迭代对象时十分有用。这是通过将作为选择器的可迭代对象取为布尔值列表来实现的。下面是它的实现方式：

>>> from itertools import compress
>>> letters = ‘ABCDEFG‘
>>> bools = [True, False, True, True, False]
>>> list(compress(letters, bools))
[‘A‘, ‘C‘, ‘D‘]

这个例子中，我们有七个字母和五个布尔值构成的列表。我们它们传入 compress 函数。compress 函数将交替查看两个可迭代对象。先检查第一个可迭代对象，然后是第二个，如果第二个的元素为 True，则保留第一个对应元素，如果为 False，则丢弃第一个对应元素。据此，再看看上面的例子，你会发现第一、第三和第四个位置元素为 True，对应的与第一个对象中的 A, C 和 D。

dropwhile(断言, 可迭代对象)

itertools 还提供了一个小清新的迭代器 dropwhile。只要过滤器判定是  True，dropwhile 迭代器就会排除这些元素。因此，在出现断言为 False 之前，你不会看到任何输出结果。这可能导致启动时间非常长，这点应当注意。

我们看一个来自于 Python 文档的例子：

>>> from itertools import dropwhile
>>> list(dropwhile(lambda x: x<5, [1,4,6,4,1]))
[6, 4, 1]

这里，我们先导入 dropwhile，然后传入一个简单的 lambda 声明的函数，如果 x 小于5，lambda 将返回 True ，否则返回 False。dropwhile 将遍历整个列表，然后将每个元素传入 lambda。如果 lambda 返回 True 则舍弃该值。一旦遇到了数字 6，lambda 将返回 False，因此，我们将保留数字 6 以及之后余下的元素。

当我了解更多之后发现，用一般的函数比 lambda 函数更为有用。所以我们尝试创建一个函数，如果参数大于 5，函数返回 True。

>>> from itertools import dropwhile
>>> def greater_than_five(x):
...     return x > 5 
... 
>>> list(dropwhile(greater_than_five, [6, 7, 8, 9, 1, 2, 3, 10]))
[1, 2, 3, 10]

这里，我们在 Python 解释器中创建了一个简单的函数作为断言或过滤器。如果我们传入的值为 True，这些值都会被舍弃。一旦我们遇到了一个小于 5 的值，那么该值，包括其后余下的全部值都将被保留，正如你在例子中看到的。

filterfalse(断言, 可迭代对象)

itertools 中的 filterfalse 函数与 dropwhile 非常类似。只是，filterfalse 返回断言为 False 的那些值，而不是舍弃断言为 True 的值。以我们上一节创建的函数为例来阐释：

>>> from itertools import filterfalse
>>> def greater_than_five(x):
...     return x > 5 
... 
>>> list(filterfalse(greater_than_five, [6, 7, 8, 9, 1, 2, 3, 10]))
[1, 2, 3]

这里，我们向 filter 传入了我们创建的函数和一个整数元素的列表。如果整数值小于 5，这个整数将保留下来，否则将被舍弃。注意到，这里的结果仅有 1, 2 和 3。与 dropwhile 不同，filterfalse将对全部的值进行条件判断。

groupby(可迭代对象, 键=None)

groupby 迭代器会从传入的可迭代对象返回连续的键和组。不借助例子可能很难理解这一点，所以我们还是看例子。将下面的代码输入解释器或者存为一个文件：

from itertools import groupby
 
vehicles = [(‘Ford‘, ‘Taurus‘), (‘Dodge‘, ‘Durango‘),
            (‘Chevrolet‘, ‘Cobalt‘), (‘Ford‘, ‘F150‘),
            (‘Dodge‘, ‘Charger‘), (‘Ford‘, ‘GT‘)]
 
sorted_vehicles = sorted(vehicles)
 
for key, group in groupby(sorted_vehicles, lambda make: make[0]):
    for make, model in group:
        print(‘{model} is made by {make}‘.format(model=model,
                                                 make=make))
    print ("***** END OF GROUP *****")

这里，我们先导入 groupby，然后创建一个元组构成的列表。接着，我们试图对数据排序使其输出时更有意义，这也让 groubby 能够正确地对元素分组。下一步，我们遍历 groupby 返回的含有键和组的迭代器。接着我们遍历组，并且打印出它的内容。如果你运行这段代码，你应该会看到下面的结果。

Cobalt is made by Chevrolet
***** END OF GROUP *****
 
Charger is made by Dodge
Durango is made by Dodge
***** END OF GROUP *****
 
F150 is made by Ford
GT is made by Ford
Taurus is made by Ford
***** END OF GROUP *****

试试改动一下上面的代码，将传入的 sorted_vehicles 替换成 vehicles。你很快就会明白为什么你要在传入 groupby 前对数据排序。

islice(可迭代变量, 起始值, 终止值[, 步长])

我们实际上在 count 一节的时候已经提到了 islice，这里我们将更深入地探讨这一函数。islice 是一个返回可迭代对象选定元素的迭代器。这种表述有点模糊。简而言之，islice 所做的是利用可迭代对象的索引实现切片，然后以迭代器的形式返回所选元素。实际上 islice 有两种实现方式，一种是 itertools.islice(iterable, stop)，还有一种更符合 Python 惯例的形式：islice(iterable, start, stop[, step])。

我们来看看第一种形式，观察它是如何工作的：

>>> from itertools import islice
>>> iterator = islice(‘123456‘, 4)
>>> next(iterator)
‘1‘
>>> next(iterator)
‘2‘
>>> next(iterator)
‘3‘
>>> next(iterator)
‘4‘
>>> next(iterator)
Traceback (most recent call last):
  Python Shell, prompt 15, line 1
builtins.StopIteration:

在以上的代码中，我们传入了一个 6 个字符组成的字符串以及终止参数 4 给 isilce。这表示 islice 返回的迭代器将包含传入字符串的前 4 个字符。为了验证，我们连续四次调用 next 函数。Python 能够智能地识别是否只有 2 个传入参数，若是，则第二个参数就取为终止参数。

我们再试试传入三个参数来验证是否可以同时传入起始值和终止值。count 工具可以帮助我们解释这一概念：

>>> from itertools import islice
>>> from itertools import count
>>> for i in islice(count(), 3, 15):
...     print(i)
... 
3
4
5
6
7
8
9
10
11
12
13
14

这里我们调用了 count，并指定 islice 从 3 开始，到 15 时结束。这跟我们用 slice 效果是一样的，不同之处在于传入的是迭代器，返回的是新的迭代器。

starmap(函数, 可迭代对象)

starmap 工具能够创建一个用传入的函数和可迭代对象计算的迭代器。如文档中所言，“map() 和 starmap() 的区别正如 function(a,b) 和 function(*c) 的区别。”

下面来看一个简单的例子：

>>> from itertools import starmap
>>> def add(a, b):
...     return a+b
... 
>>> for item in starmap(add, [(2,3), (4,5)]):
...     print(item)
... 
5
9

这里，我们先创建了一个接受两个参数的求和函数。接下来我们创建了一个 for 循环，并调用 starmap 函数，starmap 函数的第一个传入参数是我们创建的求和函数，第二个传入参数是元组构成的列表。starmap 函数接着会将每一个元组传入求和函数，然后将结果返回到迭代器，正如我们打印出的那样。

takewhile(断言, 可迭代对象)

takewhile 模块与我们之前介绍的 dropwhile 迭代器刚好相反。takewhile 所创建的迭代器，一旦可迭代对象的断言或过滤器结果为 True 就返回元素。试一试下面的例子，看看它是如何工作的

>>> from itertools import takewhile
>>> list(takewhile(lambda x: x<5, [1,4,6,4,1]))
[1, 4]

这里，运行 takewhile 时传入了一个 lambda 函数和一个列表。输出的仅是可迭代对象的前两个整数元素。因为 1 和 4 都小于 5，而 6 大于5，一旦 takewhile 遇到 6， 判定条件结果将是 False，可迭代对象余下的元素将会被忽略。

tee(可迭代对象, n=2)

tee 工具能够从一个可迭代对象创建 n 个迭代器。这意味着你能够用一个可迭代对象创建多个迭代器。下面这段代码能大体解释它是如何工作的：

>>> from itertools import tee
>>> data = ‘ABCDE‘
>>> iter1, iter2 = tee(data)
>>> for item in iter1:
...     print(item)
... 
A
B
C
D
E
>>> for item in iter2:
...     print(item)
... 
A
B
C
D
E

这里我们创建了一个 5 个字母组成的字符串，然后传入 tee。由于 tee 的默认参数是 2，我们用两个变量接收 tee 返回的两个迭代器。最后，我们分别遍历了这两个迭代器，并打印出它们的内容。正如你所见的，它们的内容是相同的。

zip_longest(*可迭代对象, 填充值=None)

zip_longest 迭代器可以用于将两个可迭代对象配对。如果可迭代对象的长度不同，也可以传入填充值。我们来看一个来自该函数文档的一个很初级的例子：

>>> from itertools import zip_longest
>>> for item in zip_longest(‘ABCD‘, ‘xy‘, fillvalue=‘BLANK‘):
...     print (item)
... 
(‘A‘, ‘x‘)
(‘B‘, ‘y‘)
(‘C‘, ‘BLANK‘)
(‘D‘, ‘BLANK‘)

这段代码里，我们先导入 zip_longest，然后传入两个需要配对的字符串。你会发现，第一个字符串有 4 个字符，而第二个只有 2 个字符。我们还传入了填充值”BLANK“。当遍历并打印的时候，会看到返回的是元组。前两个元组是分别来自两个字符串第一个和第二个字母的组合。后两个则插入了填充值。

需要注意的是，如果传入 zip_longest 的可迭代对象是无限的，那么你应该用类似 islice 的工具在外部处理函数，以便控制调用的次数。

---

 

（本博客转至伯乐在线 - 張無忌，本博客仅做学习笔记，若原作者有意见，请联系！）