Python itertools 简单介绍和运用例

前言

最近写 Python 比较多，不可避免地要处理一堆可迭代对象，发现 Python 对于迭代器/生成器的支持相较于其它语言来说是更为丰富的，所以简单记录一下 itertools 这个内置包中几个常见的函数。

文末附一个实例，是我写的一个扫雷游戏的算法，用到了文中提到的一些函数。注意，我知道它们并不是真正的函数，而是以类的形式定义。为了方便起见，本文就把它们当成函数看代。

介绍

使用迭代器/生成器的好处是节省无意义的内存消耗，就像 Java 里的 Stream API 一样（当然也有很多不同，这里不再展开讨论）。假设我们最终要对一个 iterable 进行求和或者其它 reduce 操作，那是没有必要为此创建一堆数组的。

反例如 JS 中 Array 上的 map，filter 等操作，它们都会创建一个新的数组，但是这点的性能其实不值一提，只是锦上添花而已。就像上文说的，只是节省无意义的内存消耗。

扯多了，回到我们的 itertools 上面来。

实际上还有很多，可以参考 Python 官方文档，这里只是挑了几个我比较喜欢用的。

示例

下面详细说一下其中几个，举几个例子。

repeat 可以搭配 map 一起使用，看一下代码就明白怎么用了。

>>> list(map(lambda n: pow(n, 3), range(5)))
[0, 1, 8, 27, 64]
>>> list(map(pow, range(5), repeat(3)))
[0, 1, 8, 27, 64]

因为 map 在任何一个可迭代对象结束之后就会结束，所以我们直接让它无限重复 3 作为 pow 的第二个参数，就达到了我们求立方的目的。

cycle 与 islice 一起使用，重复从一个可迭代对象中取出特定个元素。

>>> ''.join(islice(cycle('ABCD'), 7))
'ABCDABC'

本来我还想说 count 和 islice 一起用，但我又想到这不就是 range 么，不说了喵。

tee 的实现原理挺有意思的，我这里把官方文档里的等价代码粘贴过来。

def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # when the local deque is empty
                try:
                    newval = next(it)   # fetch a new value and
                except StopIteration:
                    return
                for d in deques:        # load it to all the deques
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

它使用了先进先出的队列来保存迭代器返回的数据，实现还是比较巧妙的，下面贴上使用方法。

>>> a, b = tee(range(10))
>>> for i in a:
...     print(i)
...
>>> for j in b:
...     print(j)

但其实当数据比较大的时候，它还是会消耗很多内存，所以慎重使用。

product 可以展平嵌套 for 循环，以下两种写法结果相同。我这里就不把输出粘贴过来了，有点长。

>>> for x in range(2):
...    for y in range(3):
...        print(x, y)
...
>>> for x, y in product(range(2), range(3)):
...    print(x, y)

下面是计算扫雷游戏中每个格子周围有几个雷的算法。

for x, y in product(range(width), range(height)):
    self[x, y].count = sum(
        self[x + dx, y + dy].is_mine
        for dx, dy in product((-1, 0, 1), repeat=2)
    )

我只把和本节相关的代码贴了出来，下面是一些解释。

• 这是一个实例方法，它所在的类实现了__getitem__，当坐标无效时会返回一个所有的属性都是 0 或 False 的对象，避免进行判空。
• 格子对象的 is_mine 是一个布尔值，它可以被当成 int 参与运算。
• 这段代码也会把当前格子是否是雷给计算上，它对游戏是无关紧要的，因为当这个格子是雷的时候，计数也不会被显示出来。
• 如果你对实现感兴趣，可以前往这个仓库中 board.py 和 block.py。我在写完这个项目的时候才知道有人已经实现了，所以我就直接 Archive 了，也算间接导致了这篇文章的出现。

就这样，白白喵。