itertools 提供一组按需产生数据的迭代器。排列、组合和笛卡尔积已经在暴力验证主文章中使用;本文补充相邻元素、前缀聚合、连接序列和不等长并行遍历。
这些函数返回迭代器,不会自动创建完整列表。需要调试输出或多次遍历时再使用 list(...) 保存结果。
pairwise:枚举相邻元素
1from itertools import pairwise
2
3a = [3, 5, 8, 13]
4
5assert list(pairwise(a)) == [(3, 5), (5, 8), (8, 13)]
检查严格递增:
1from itertools import pairwise
2
3
4def is_strictly_increasing(a):
5 return all(x < y for x, y in pairwise(a))
6
7
8assert is_strictly_increasing([1, 3, 8])
9assert not is_strictly_increasing([1, 3, 3])
10assert is_strictly_increasing([])
pairwise 从 Python 3.10 开始提供。旧版本使用:
1a = [3, 5, 8, 13]
2
3assert list(zip(a, a[1:])) == [(3, 5), (5, 8), (8, 13)]
切片会创建一个新列表,验证用小数据通常无需担心这点。
accumulate:前缀聚合
默认情况下,accumulate 产生前缀和:
1from itertools import accumulate
2
3a = [3, -2, 5, -1]
4prefix = list(accumulate(a))
5
6assert prefix == [3, 1, 6, 5]
添加 initial=0 可以得到竞赛中常用的长度为 的前缀和:
1from itertools import accumulate
2
3a = [3, -2, 5, -1]
4prefix = list(accumulate(a, initial=0))
5
6assert prefix == [0, 3, 1, 6, 5]
7
8left, right = 1, 3
9assert sum(a[left:right]) == prefix[right] - prefix[left] == 3
也可以传入其他二元函数:
1from itertools import accumulate
2from operator import mul
3
4assert list(accumulate([2, 3, 4], mul)) == [2, 6, 24]
5assert list(accumulate([3, 1, 5, 2], max)) == [3, 3, 5, 5]
initial 从 Python 3.8 开始提供。旧版本可以写 [0] + list(accumulate(a))。
chain:连接多个序列
1from itertools import chain
2
3left = [1, 2]
4middle = (3, 4)
5right = {5, 6}
6
7combined = list(chain(left, middle, sorted(right)))
8
9assert combined == [1, 2, 3, 4, 5, 6]
它相当于依次遍历多个容器,不创建中间拼接列表。容器本身放在一个列表中时,使用 chain.from_iterable:
1from itertools import chain
2
3groups = [[1, 2], [], [3], [4, 5]]
4
5assert list(chain.from_iterable(groups)) == [1, 2, 3, 4, 5]
列表嵌套只有一层时,这是一种直接的展平方法。
repeat:重复同一个值
1from itertools import repeat
2
3assert list(repeat("x", 4)) == ["x", "x", "x", "x"]
它可以为 map 的某个参数提供固定值。例如计算每个数的平方:
1from itertools import repeat
2
3bases = [2, 3, 4]
4squares = list(map(pow, bases, repeat(2)))
5
6assert squares == [4, 9, 16]
这里 map 每次分别从 bases 和 repeat(2) 取一个参数,调用 pow(base, 2)。
不提供次数时,repeat(value) 是无限迭代器,必须由较短的输入、islice 或短路条件限制。
zip_longest:并行遍历不同长度序列
普通 zip 在最短序列结束时停止:
1assert list(zip([1, 2, 3], [10, 20])) == [(1, 10), (2, 20)]
zip_longest 会继续到最长序列结束,并为缺失位置填默认值:
1from itertools import zip_longest
2
3pairs = list(zip_longest([1, 2, 3], [10, 20], fillvalue=0))
4
5assert pairs == [(1, 10), (2, 20), (3, 0)]
例如逐位相加两个低位在前的数字数组:
1from itertools import zip_longest
2
3
4def add_digits_without_carry(a, b):
5 return [
6 x + y
7 for x, y in zip_longest(a, b, fillvalue=0)
8 ]
9
10
11assert add_digits_without_carry([3, 2, 1], [7, 6]) == [10, 8, 1]
如果长度不同本身就是错误,不应使用 zip_longest 掩盖问题,应先显式断言长度相等。
islice:截取迭代器
切片语法不能直接用于生成器,islice 可以按位置截取:
1from itertools import islice
2
3squares = (x * x for x in range(100))
4
5assert list(islice(squares, 3, 7)) == [9, 16, 25, 36]
它不会倒序,也不支持负下标,因为普通迭代器只能向前产生元素。
与排列组合工具的关系
| 任务 | 工具 |
|---|---|
| 相邻元素 | pairwise |
| 前缀和、前缀最大值 | accumulate |
| 顺序连接多个序列 | chain |
| 提供重复参数 | repeat |
| 不同长度并行遍历 | zip_longest |
| 截取迭代器 | islice |
| 每个位置选择一种状态 | product |
| 枚举顺序 | permutations |
| 选择若干元素 | combinations |
最后三项直接枚举指数级或阶乘级状态空间,详见用 Python 快速编写算法暴力验证程序。生成器的一次性消费、短路与内存特点见Python 生成器表达式。