Python 模块：itertools

在某些情况下，我们通常需要对序列进行一些复杂的操作，比如从序列中选出一部分元素做排列，组合，笛卡尔积等。如果自己实现这个操作未免太繁琐了，而且还会占用大量的空间，这个时候我们可以求助于 Python 模块——itertools。这个模块总共有 3 部分——无穷迭代器，根据最短输入序列的长度停止的迭代器，排列组合迭代器。

无穷迭代器

常用的无穷迭代器总共有 3 个——count，cycle，repeat，我们先来看一下 count。

count

love

count 顾名思义——数数，这个函数有两个默认参数，第 1 个参数表示从几开始数（默认值 0），第 2 个参数表示每次数的步长（默认值 1）。

import itertools
count = itertools.count()
for c in count:
    print(c)

运行结果输出 0 1 2 3 4 5……它会一直数下去，根本停不下来。

cycle

love

cycle 顾名思义——周期，这个函数有一个必选参数，这个必选参数必须是可迭代的。

import itertools
cycle = itertools.cycle([1, 2, 3])
for c in cycle:
    print(c)

运行结果输出 1 2 3 1 2 3……一直周而复始，永不停息。

repeat

love

repeat 顾名思义——重复，这个函数有一个必选参数，这个必选参数类型是任意的。

import itertools
cycle = itertools.repeat(None)
for c in cycle:
    print(c)

运行结果会一直重复输出 None，也不会停。

根据最短输入序列的长度停止的迭代器

这样的迭代器就比较多了，常用的有 accumulate，chain，chain.from_iterable，compress，dropwhile，filterfalse，islice，starmap，takewhile，ziplongest。

accumulate

love

accumulate 顾名思义——累计，传入一个必选参数和一个可选参数，其中必选参数是一个可迭代对象，可选参数是一个函数（默认值 operator.add）。

import itertools
a = itertools.accumulate([1, 2, 3])
for a0 in a:
    print(a0)

运行结果如图所示。

1 前面没有东西，所以还是 1，2 前面有一个 1，所以 2+1 = 3，3 前面有 2 和 1，所以 3+2+1 = 6。如果我不想做累加，想做累乘怎么办？很简单，指定第二个参数就可以了，就像这样：itertools.accumulate([1, 2, 3], operator.mul)。

chain

love

chain 顾名思义——链条，主要用来把多个序列连在一起做迭代。有些人会认为把多个序列连在一起完全使用操作符 + 就可以实现了，确实可以，但是这特别消耗内存资源，比如下面这种情况。

import itertools
a = [1, 2, 3]
b = [4, 5, 6]
a = a+b

刚开始是两个长度为 3 的列表，当执行完 a = a+b 这一行时，有些人会认为这就仅仅是一个长度为 6 的列表，其实并不是，目前变量 b 所指向的内存并没有被 free，因此内存中还是有 2 个列表，一个长度为 3，一个长度为 6，很明显这比做 + 之前占用的内存要多。为了避免这种情况，我们可以使用 chain。chain 函数参数个数不定，但要求每一个参数都是可迭代对象。

import itertools
chain = itertools.chain([1, 2, 3], [4, 5, 6])
for c in chain:
    print(c)

运行结果如图所示。

chain.from_iterable

love

和 chain 一样，只不过这次不是传入多个参数，而是直接传入一个可迭代对象，如下所示。

import itertools
chain = itertools.chain.from_iterable([[1, 2, 3], [4, 5, 6]])
for c in chain:
    print(c)

运行结果同上。

compress

love

compress 是用来做筛选的，两个必选参数，第一个是待筛选的序列 data，第二个是选择器列表 selectors，如果 selectors[i] 为真，就返回 data[i]。注意：如果 data 和 selectors 长度不相等，以最短的那个为基准！

import itertools
compress = itertools.compress('ABCDEF', [1, 0, 1, 0, 1, 1])
for c in compress:
    print(c)

运行结果如图所示。

dropwhile 和 takewhile

love

这个函数有两个参数，第一个参数是一个函数 pred，第二个参数是一个序列 seq。如果 pred(seq[i]) 为真，直接丢弃，如果 pred(seq[i]) 为假，就不丢弃，注意：当 pred(seq[i]) 为假时，即使 pred(seq[j]) 为真也不会被丢弃（j > i）。

import itertools
d = itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])
for d0 in d:
    print(d0)

运行结果如图所示。

很明显，前两个元素被丢弃，后三个元素被保留在迭代器中被迭代输出了。顺便说一下 takewhile，它和 dropwhile 相反，如果这段代码把 dropwhile 改成 takewhile，输出结果就是 1 4，也就是前两个元素。

starmap

love

starmap 是 map 的一个升级版本，比如有这么一个要求，一个列表 a = [(2, 5), (3, 2), (10, 3)]，计算 pow(2, 5)，pow(3, 2)，pow(10, 3)。用 map 实现确实可以，但是 map 每次计算都是不带解包的，所以要想使用 map 必须要多一步解包的操作，显得有些繁琐。而 starmap 给你封装好了，直接用即可。

import itertools
starmap = itertools.starmap(pow, [(2, 5), (3, 2), (10, 3)])
for s in starmap:
    print(s)

运行结果如图所示。

zip_longest

love

zip_longest 是相对于 zip 而言的，这两个函数的第一个参数一样，都是多个可迭代对象。执行的时候都是从每个参数中收集对应的元素，先收集第一个参数，第二个参数……的第一个元素，然后收集第一个参数，第二个参数……的第二个元素，以此类推。那么 zip 和 zip_longest 有什么区别呢？当每个参数的长度不一样时，zip 以最短的为基准，而 zip_longest 以最长的为基准，那么 zip_longest 有缺失怎么办，其实这个函数还有一个默认参数 fillvalue（默认值 None），这个参数就是用来决定有缺失时填充什么。

import itertools
a = itertools.zip_longest('ABCD', 'xy', fillvalue='-')
for a0 in a:
    print(a0)

运行结果如图所示。

排列组合迭代器

排列组合迭代器有 product，permutations，combinations，combinations_with_replacement。

product

love

product 主要用来计算多个序列的笛卡尔积，传入一个可变长参数和一个默认参数 repeat（默认值 1）。这个可变长参数要求每一个都是可迭代对象。

import itertools
x = [1, 2, 3]
y = [4, 5, 6]
a = itertools.product(x, y)
for a0 in a:
    print(a0)

运行结果如图所示。

如果有两个一样的序列 x 做笛卡尔积，我是不是要写成 product(x, x) 呢？完全没有必要，只要使用 repeat 参数就可以了，写成 product(x, repeat=2) 就可以了。

permutations

love

这个函数就是用来做排列的，比如我输入 (1, 2, 3)，输出 (1, 2, 3)，(1, 3, 2)，(2, 1, 3)，(2, 3, 1)，(3, 1, 2)，(3, 2, 1)，我就可以使用这个函数，如下所示。

import itertools
a = itertools.permutations([1, 2, 3])
for a0 in a:
    print(a0)

运行结果如图所示。

那么如果我想输出 (1, 2)，(1, 3)，(2, 1)，(2, 3)，(3, 1)，(3, 2) 该怎么办呢？其实很简单，permutations([1, 2, 3], 2) 就可以了。

combinations 和 combinations_with_replacement

love

这两个函数都是是用来做组合的，从一个序列中选取几个元素，迭代所有的可能。当然这两个是有区别的，其中 combinations 是不能重复选择，而 combinations_with_replacement 是可以重复选择的。

import itertools
print('combinations:')
a = itertools.combinations([1, 2, 3, 4, 5], 2)
for a0 in a:
    print(a0)
print('combinations_with_replacement:')
a = itertools.combinations_with_replacement([1, 2, 3, 4, 5], 2)
for a0 in a:
    print(a0)

运行结果如图所示。