Python函数式编程09 itertools模块：更多迭代器

itertools 模块中包含大量迭代器函数，是处理数据的非常好用的工具。

生成器是迭代器的一个部分，迭代器和生成器的用途是一致的，只是迭代器并不都是使用 yield 这种关键字的方式定义。可以用生成器的概念来理解它。

无限迭代器

以下迭代器如果在迭代时不使用 break 终止或遇到错误，那么它会一直执行下去。

count：计数迭代器

内置的 range() 函数需要定义范围的上界，而下限和步长值可选。count(start=0, step=1) 函数与之相反，需要给出可选的起始值和步长，无须定义上界。在停止之前，该生成器会一直执行下去。

count() 函数的参数可以是浮点数，但是由于进制带来的舍入误差会随着程序的运行逐渐累积，因此尽量在较大的计数过程中慎用浮点数计数，或者使用整型计数再转换为浮点数的方式。

count(10)        # => 10 11 12 13 14 ...
count(2.5, 0.5)  # => 2.5 3.0 3.5 ...

cycle：循环迭代器

cycle() 函数重复循环一组值，可用它循环数据集标识符对数据集进行分组。

cycle('ABC') # => A B C A B C A B C ...

repeat：重复迭代器

repeat(elem, n=None) 用于重复单个值，不过也可以通过第二个参数将其限制为有限的迭代器。

repeat(10, 3) # => 10 10 10

有限迭代器

有限迭代器更多，也更常用。它们具体的长度一般取决于给定序列的长度或特征。

accumulate：归约迭代器

accumulate() 函数基于给定的函数返回一个可迭代对象，将一系列归约值汇总在一起，遍历迭代器得出当前汇总值。以下给出了一个这样的示例：

accumulate([1, 2, 3, 4, 5]) # => 1 3 6 10 15

accumulate() 函数的第二个参数 func 的默认行为是加法，也可以改成其它的归约函数。

chain：组合迭代器

chain() 函数将多个迭代器组合为单个迭代器，可以用于一次处理多个集合。

chain(range(5), range(10, 20, 2)) # => 0 1 2 3 4 10 12 14 16 18

如果多个迭代器被放在一个可迭代对象内，那么还可以使用 chain.from_iterable() 类方法取代以上这种可变参数的形式。

chain.from_iterable(['ABC', 'DEF']) # => A B C D E F

groupby：分组迭代器

groupby(iterable, key=None) 通过对每个元素应用 key 函数求值，并根据求值结果将一个迭代器切分为多个小迭代器，具体的规则为：如果后一个元素的 key 返回值等于前一个元素的 key 返回值，会将这两个元素放在同一个分组中；如果与前一个元素的 key 值不同，则当前分组结束，将当前元素放到新的分组中。

该生成器每一次迭代的结果都是一个 (key, group) 形式的元组，其中第二个元素是生成器。

以下是一个简单的示例，将一个排序好的列表按十位分组：

from random import randint
target = filter(lambda _: randint(0, 1), range(50))
for i in groupby(target, key=lambda x: x // 10):
    print(i[0], list(i[1]))

结果为：

$ python -u demo.py
[2, 4, 6]
[10, 11, 13, 16, 19]
[23, 26, 27, 29]
[30, 31, 32, 35, 37, 38]
[40, 41, 46, 47, 48]

因此，groupby() 函数的输入列表必须是排序好的，以确保分在一组中的元素是相邻的。

compress：选择迭代器

内置的 filter() 根据函数求值结果为真或假决定是否保留元素，而 compress(data, selectors) 则根据 selectors 可迭代对象给出的布尔值决定是否保留同一位置处的元素。

compress('ABCDEF', [1, 0, 1, 0, 1, 0]) # => A C E

islice：切片迭代器

islice(iterable, start=0, stop, step=1) 函数可以对迭代器实现切片操作，就像 slice 类对序列实现切片操作一样。

islice('ABCDEFG', 2)           # => C D E F G
islice('ABCDEFG', 1, None, 2)  # => B D F

dropwhile和takewhile：过滤状态迭代器

dropwhile(predicate, iterable) 和 takewhile(predicate, iterable) 都是过滤函数，其用法是：从 predicate 函数给定的一种布尔模式开始，当布尔状态变换后，则函数的过滤规则也发生变化。dropwhile() 函数开始采用拒绝模式，当谓词函数变为 False 时切换并一直保持为通过模式。takewhile() 则从通过模式开始，当谓词函数变为 False 时切换并一直保持为拒绝模式。

dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])  # => 6 4 1
takewhile(lambda x: x < 5, [1, 4, 6, 4, 1])  # => 1 4

这两个函数的一种常见用法是过滤掉文件的头部和尾部的一些无用信息。例如，假设有如下附加了额外头部信息的 CSV 文件：

Author: Hello
Date: 2022-05-31
sepal_length,sepal_width,petal_length,petal_width,species
5.1,3.5,1.4,0.2,Iris-setosa
4.9,3.0,1.4,0.2,Iris-setosa
6.4,3.2,4.5,1.5,Iris-versicolor
6.9,3.1,4.9,1.5,Iris-versicolor

那么可以使用如下方式处理它：

with open('header.csv', 'r') as f:
    data = dropwhile(lambda row: not row[0].isdigit(), f)
    for i in map(lambda row: row.replace('\n', '').split(','), data):
        print(i)

该程序检查 CSV 文件的每一行（注意文件对象也是可迭代对象，每次迭代结果为文件中的一行），判断每一行的首字符是不是数字，并将头部非数字开头的行丢弃。其结果为：

$ python -u demo.py
['5.1', '3.5', '1.4', '0.2', 'Iris-setosa']
['4.9', '3.0', '1.4', '0.2', 'Iris-setosa']
['6.4', '3.2', '4.5', '1.5', 'Iris-versicolor']
['6.9', '3.1', '4.9', '1.5', 'Iris-versicolor']

改进的迭代器

以下迭代器是对内置迭代器的补充或改进。

zip_longest：不截短的zip

在第 4 节介绍了 zip() 迭代器，它接收多个序列，返回结果的长度是其中最短序列的长度。而 zip_longest(*iterables, fillvalue=None) 则为较短的序列填充 fillvalue 参数提供的默认值，直到遍历完最长的序列。

zip_longest('ABCD', 'xy', fillvalue='*') # => Ax By C* D*

filterfalse：filter的反相过滤

内置的 filter() 迭代器会保留所有求值结果为 True 的数据；而 filterfalse() 则与之相反，它会保留所有求值结果为 False 的数据。

filter(lambda x: x % 2, range(10))       # => 1 3 5 7 9
filterfalse(lambda x: x % 2, range(10))  # => 0 2 4 6 8

通过组合这两个迭代器，可以方便地将输入数据分为“保留”和“舍弃”两组。

starmap：平铺的map

内置的 map() 迭代器可以接收多个可迭代对象，它将每一个可迭代对象中的元素组合为所需的参数；而 starmap() 迭代器只接收一个可迭代对象，它将其中元素展开后变成所需的参数。以下给出了一个这样的示例：

starmap(pow, [(2,5), (3,2), (10,3)]) # => 32 9 1000

因此，使用 zip() 迭代器可以非常简单地将 map() 变成 starmap() ：

def starmap(function, *iterables):
    return map(function, zip(iterables))

tee：克隆迭代器

tee(iterable, n=2) 函数可以突破迭代器只能使用一次的限制，将可迭代对象拷贝 n 份，并将拷贝结果以元组的形式返回。

不过该函数在使用时有一些限制：它的实现实际上有将迭代器转化为序列的过程，因此在处理大型数据集且拷贝次数较少时的效果往往不佳。

组合迭代器

接下来的迭代器与排列组合有关，因此它们往往需要输入多个序列。

product：笛卡儿积

product(*iterables, repeat=1) 用于得到笛卡儿积，即基于一组集合生成所有可能的元素组合。repeat 参数可以将提供的集合重复若干次参与组合。例如：

product(range(3), range(3))  # => 00 01 02 10 11 12 20 21 22
product('AB', repeat=2)      # => AA AB BA BB

该函数的作用结果类似嵌套的 for 循环。以下再次给出了一个示例：

>>> from random import choice
>>> cards = list(product(range(1, 14), '♣♦♥♠'))
>>> choice(cards)
(2, '♦')
>>> choice(cards)
(9, '♥')

排列组合相关迭代器

permutations(iterable, r=None) 函数会排列集合中所有元素；如果给出参数 r ，则它会限制参与排列元素的个数，例如：

permutations(range(3))   # => 012 021 102 120 201 210
permutations('ABCD', 2)  # => AB AC AD BA BC BD CA CB CD DA DB DC

由于 \\( n \\) 个元素的排列结果有 \\( P_n^r = \frac{n!}{(n-r)!} \\) 个，因此随着元素个数的上升，排列结果的长度急剧上升。

与之类似的函数是 combinations(iterable, r) ，它会以忽略顺序的形式组合集合中所有元素。

combinations('ABCD', 2)    # => AB AC AD BC BD CD
combinations(range(4), 3)  # => 012 013 023 123

还有一个类似的迭代器 combinations_with_replacement(iterable, r) ，其区别是该迭代器在组合时允许元素重复：

combinations('ABC', 2)                   # => AB AC BC
combinations_with_replacement('ABC', 2)  # => AA AB AC BB BC CC

这些迭代器通常用于处理排列组合相关的数学问题中。

使用这些迭代器可以代替许多算法中出现的简单循环，并使执行过程更高效。itertools 模块中还给出了更多范例，将基本的迭代器组合成更高级、有用的迭代器，可以高效地处理复杂的数据，参见 https://docs.python.org/3/library/itertools.html#itertools-recipes 。

参考资料

https://docs.python.org/3/library/itertools.html

itertools 模块的官方文档。本文中的许多示例代码都来自该文档中。

文章中使用到的 CSV 数据来自 kaggle 。