Python提供了一个非常棒的模块用于创建自定义的迭代器,这个模块就是 itertools。itertools 提供的工具相当高效且节省内存。使用这些工具,你将能够创建自己定制的迭代器用于高效率的循环。这一章,我们将一起看一看这些工具的应用实例以便理解并应用到自己的编程中去。
让我们先从几个无限迭代器的例子开始吧!
无限迭代器
itertools 包自带了三个可以无限迭代的迭代器。这意味着,当你使用他们时,你要知道你需要的到底是最终会停止的迭代器,还是需要无限地迭代下去。
这些无限迭代器在生成数字或者在长度未知的可迭代对象(iterables)中循环时相当有用。下面我们开始认识这些有趣的可迭代对象!
count(初值=0, 步长=1)
count 迭代器会返回从传入的起始参数开始的均匀间隔的数值。count 也可以接收指定的步长参数。我们来看一个简单的例子:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
>>> 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。使用方法如下:
1 2 3 4 5 6 7 8 9 |
>>> 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 个字母的字符串,看看将会发生什么:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
>>> 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 创建的迭代器进行循环:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
>>> 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 不会在多个值之间循环。我们看一个简单的例子:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
>>> 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 迭代器将返回累计求和结果,或者传入两个参数的话,由传入的函数累积计算的结果。默认设定为相加,我们赶快试一试吧:
1 2 3 |
>> 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 模块,然后添加进去:
1 2 3 |
>>> 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 迭代器能够将多个可迭代对象合并成一个更长的可迭代对象。实际上,我参与的一个项目中最近就需要这一功能。我有一个列表,里面已经包含一些元素,接着想把另外两个列表添加到最初那个列表中。注意,我们想添加的是两个列表的元素。最初,我是这样做的:
1 2 3 4 5 6 |
>>> 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 来合并这些列表:
1 2 3 4 |
>>> 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,也有其他方法能够实现这一要求。你可以这样做:
1 2 3 4 |
>>> 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 有些细微的差别。不同于直接传入一系列可迭代对象,你必须传入一个嵌套的列表。我们这就来看一看