【算法】【Python】itertools包的妙用 - AI科研编程读书笔记

本文是一个全面总结，关于itertools包在Python算法竞赛中的技巧性使用方法和典型应用场景。itertools是Python标准库中一个功能强大且高效的迭代器模块，特别适合在算法竞赛中处理组合、排列、枚举、笛卡尔积、累加、分组等问题，能大幅度简化代码量，提高效率。

一、常用函数与技巧总结

1. permutations：排列（考虑顺序）

用法： 生成列表或字符串中所有长度为r的排列（默认r为原长度），考虑顺序。

常见场景：

全排列问题（如 TSP 问题，字典序枚举）
判断是否存在某种排列满足某个条件

示例：

from itertools import permutations
for p in permutations([1, 2, 3]):
    # 每个p是一个元组
    pass

2. combinations：组合（不考虑顺序）

用法： 从iterable中取r个元素的所有组合，顺序不重要。

常见场景：

子集枚举（选出若干元素）
子集和问题、背包问题中的状态遍历
枚举两两配对等问题

from itertools import combinations
for c in combinations([1, 2, 3], 2):
    pass

3. combinations_with_replacement：允许重复的组合

用法： 生成r长度的组合，允许元素重复选取。

常见场景：

完全背包模型
多重组合问题

from itertools import combinations_with_replacement
combinations_with_replacement([1, 2], 2)  # [(1,1), (1,2), (2,2)]

4. product：笛卡尔积（排列，允许重复）

用法： 对多个可迭代对象取笛卡尔积，可以控制重复的深度。

常见场景：

枚举所有可能的密码组合（如4位数字密码）
多重嵌套循环的简化
枚举多个独立变量的所有组合

from itertools import product
product('01', repeat=3)  # 等价于三层for循环

5. accumulate：前缀和 / 累积

用法： 生成输入序列的前缀和，或传入二元操作函数自定义累积方式。

常见场景：

滑动窗口前缀和优化
区间和问题
动态规划转移中需要连续累积的地方

from itertools import accumulate
accumulate([1, 2, 3, 4])  # [1, 3, 6, 10]

也可以自定义操作，如累乘：

import operator
accumulate([1, 2, 3, 4], func=operator.mul)  # [1, 2, 6, 24]

# 求op列表的前缀和并返回给op，注意这里要用list转换类型
op = list(accumulate(op))

6. groupby：按条件分组

用法： 将连续相同的元素按键值函数分组（必须先排序）

常见场景：

统计出现次数
快速构造分组结构（如按首字母分组字符串）

from itertools import groupby
groupby(sorted('aabbcc'))

注意：只有相邻的元素才会被分为一组，因此在使用前需要对数据排序。

7. count：无限计数器

用法： 创建一个无限增长的迭代器，常与 zip 或 map 连用

常见场景：

自定义索引
用于生成测试数据、坐标等

from itertools import count
zip(count(1), ['a', 'b', 'c'])  # [(1, 'a'), (2, 'b'), (3, 'c')]

8. cycle：无限循环一个序列

常见场景：

模拟循环队列
周期性任务调度
快速构造轮转逻辑

from itertools import cycle
c = cycle([1, 2, 3])  # 1,2,3,1,2,3,...

9. repeat：重复元素

常见场景：

初始化填充数组
配合 map 进行参数复用

from itertools import repeat
list(repeat(10, 3))  # [10, 10, 10]

二、算法竞赛中高效使用场景

1. 枚举排列组合类搜索问题

如在搜索剪枝前先枚举所有可选序列，permutations 和 combinations 能替代繁琐的递归写法，并且速度不低，适合小规模搜索（n <= 8）。

2. 多重循环的替代

如三层嵌套循环：

for i in range(2):
    for j in range(2):
        for k in range(2):
            pass

可替换为：

from itertools import product
for i, j, k in product(range(2), repeat=3):
    pass

可读性更高，可灵活调整层数。

3. 状态压缩 + 子集遍历优化

有些子集可以用combinations高效枚举，比如：

枚举所有选k个物品的组合状态
枚举所有含特定条件的子集（如所有和为偶数的子集）

4. 前缀和/积/最大值等累积操作优化

accumulate比手写前缀和更简洁，适合竞赛时快写代码。

支持自定义函数，非常适合构造区间的累乘、最大值等前缀状态。

5. 分组分析优化

使用groupby代替手动维护哈希分组，适合已排序的数据聚类分析。

三、使用技巧总结

itertools的函数返回的是迭代器，通常使用list()包裹转为列表再操作。
结合map()、filter()可以进行链式高效处理。
对于product()和permutations()等函数，数据量大时要注意时间和空间复杂度，itertools也不是万能的😂。
可以与functools中的reduce()等函数结合使用构造更复杂的状态处理。
在竞赛中，需要权衡itertools的简洁性与暴力枚举的复杂度，避免在高数据量时引发TLE。

若无特殊说明，文章版权归作者所有，请勿转载至任何平台。

THE END

Python 算法