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

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

介绍

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

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

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

名称 说明 示例
count(start=0, step=1) 参数可以为 intfloat,类似于 range 但它是无限的。 count(10, 2) --> 10 12 14 ...
cycle(p) 重复遍历一个可迭代对象。 cycle('ABC') --> A B C A ...
repeat(elem, times=None) 重复同一个元素,如果不指定次数 times 就是无限重复。 repeat(10, 3) --> 10 10 10
chain(p, q, ...) 把所有的参数连接成一个可迭代对象。 chain('ABC', 'DEF') -> A B C D E F
chain.from_iterable(p) 相当于把参数解包传给 chain 函数。 chain(['ABC', 'DEF']) --> A B C D E F
comparess(data, selectors) selectors 是否为 truthy value(这是 JS 的说法,就是它们是否被当作 true)来筛选数据。 compress('ABC', [0, 1, 1]) --> B C
product(p, q, ..., repeat=1) 相当于嵌套 for 循环。 product('AB', 'CD') --> AC, AD, BC, BD
pairwise(iterable) 成对迭代一个可迭代对象。 pairwise('ABCDE') --> AB BC CD DE
takewhile(predicate, iterable) 直到 predicate(i) 返回一个 falsy value takewhile(lambda x: x<5, [1, 4, 6, 1]) -> 1 4
dropwhile(predicate, iterable) predicate(i) 返回一个 falsy value 时开始。 dropwhile(lambda x: x<5, [1, 4, 6, 1]) --> 6 1
islice(iterable, stop), islice(iterable, start, stop, step=1) 针对一个可迭代对象的切片操作。 见下文
tee(iterable, n=2) 将一个可迭代对象分裂成多个,不是线程安全的。 见下文

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

示例

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

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

    1
    2
    3
    4
    >>> 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 的第二个参数,就达到了我们求立方的目的。

  • cycleislice 一起使用,重复从一个可迭代对象中取出特定个元素。

    1
    2
    >>> ''.join(islice(cycle('ABCD'), 7))
    'ABCDABC'
  • 本来我还想说 countislice 一起用,但我又想到这不就是 range 么,不说了喵。

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

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    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)

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

    1
    2
    3
    4
    5
    6
    >>> a, b = tee(range(10))
    >>> for i in a:
    ... print(i)
    ...
    >>> for j in b:
    ... print(j)

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

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

    1
    2
    3
    4
    5
    6
    >>> 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)
  • dropwhile 可以忽略前几个不符条件的元素,takewhile 可以只保留前几个符合条件的元素。用处还是挺大的。

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

    1
    2
    3
    4
    5
    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__,当坐标无效时会返回一个所有的属性都是 0False 的对象,避免进行判空。

    • 格子对象的 is_mine 是一个布尔值,它可以被当成 int 参与运算。

    • 这段代码也会把当前格子是否是雷给计算上,它对游戏是无关紧要的,因为当这个格子是雷的时候,计数也不会被显示出来。

就这样,白白喵。