python标准库--itertools - 迭代器工具在算法比赛的应用

目录

一、基本用法与参数

1. 排列组合工具

permutations(iterable, r=None)

combinations(iterable, r)

product(*iterables, repeat=1)

2. 无限迭代器

count(start=0, step=1)

cycle(iterable)

3. 迭代器操作

chain(*iterables)

compress(data, selectors)

groupby(iterable, key=None)

4. 累积与切片

accumulate(iterable, func=operator.add)

islice(iterable, start, stop[, step])

5. 实战参数技巧

动态生成排列组合

批量处理数据

二、应用场景 

1. 排列组合工具

场景：枚举所有可能性

2. 无限迭代器

场景：周期性操作或生成序列

3. 迭代器拼接与分组

场景：数据预处理

4. 组合生成器

场景：高效生成特定模式

5. 算法竞赛实战案例

1.计算所有子数组的和（前缀和优化）

核心思路：前缀和数组

完整代码

2.生成格雷码（Gray Code）

核心思路：位运算公式

完整代码

---

一、基本用法与参数

1. 排列组合工具

permutations(iterable, r=None)

• 参数：
  • iterable：可迭代对象（如列表、字符串）。
  • r：每个排列的长度（默认等于 len(iterable)）。
• 返回：长度为 r 的所有排列（元组形式）。
• 应用：枚举所有可能的顺序（如 TSP 问题）。

from itertools import permutationslist(permutations('ABC', 2))  # 输出：[('A', 'B'), ('A', 'C'), ('B', 'A'), ...]

combinations(iterable, r)

• 参数：同 permutations，但 r 必须指定。
• 返回：长度为 r 的所有组合（不重复）。
• 应用：子集枚举（如背包问题候选解）。

from itertools import combinationslist(combinations([1, 2, 3], 2))  # 输出：[(1, 2), (1, 3), (2, 3)]

product(*iterables, repeat=1)

• 参数：
  • *iterables：多个可迭代对象。
  • repeat：重复次数（笛卡尔积的维度）。
• 返回：多个迭代器的笛卡尔积。
• 应用：多维状态枚举（如棋盘坐标组合）。

from itertools import productlist(product('AB', [1, 2], repeat=1))  # 输出：[('A', 1), ('A', 2), ('B', 1), ('B', 2)]

2. 无限迭代器

count(start=0, step=1)

• 参数：
  • start：起始值（默认 0）。
  • step：步长（默认 1）。
• 返回：无限递增的迭代器。
• 应用：生成序号（如测试用例 ID）。

from itertools import countcounter = count(10, 2)  # 10, 12, 14, ...
next(counter)  # 输出：10

cycle(iterable)

• 参数：可迭代对象。
• 返回：循环迭代器（无限重复原序列）。
• 应用：周期性操作（如轮流处理任务）。

from itertools import cyclecycler = cycle('ABC')  # A, B, C, A, B, ...
next(cycler)  # 输出：A

3. 迭代器操作

chain(*iterables)

• 参数：多个可迭代对象。
• 返回：拼接后的迭代器。
• 应用：合并多个输入源。

from itertools import chainlist(chain('ABC', [1, 2]))  # 输出：['A', 'B', 'C', 1, 2]

compress(data, selectors)

• 参数：
  • data：数据迭代器。
  • selectors：布尔掩码迭代器。
• 返回：保留 selectors 中为 True 的对应元素。
• 应用：快速过滤无效数据。

from itertools import compresslist(compress('ABCDEF', [1, 0, 1, 0, 1, 1]))  # 输出：['A', 'C', 'E', 'F']

groupby(iterable, key=None)

• 参数：
  • iterable：可迭代对象。
  • key：分组函数（默认按元素自身分组）。
• 返回：(key, group) 对的迭代器（需先排序）。
• 应用：数据分类（如统计字符串中连续相同字符的长度）。

from itertools import groupbydata = [('A', 1), ('A', 2), ('B', 3)]
for key, group in groupby(data, lambda x: x[0]):print(key, list(group))
# 输出：
# A [('A', 1), ('A', 2)]
# B [('B', 3)]

4. 累积与切片

accumulate(iterable, func=operator.add)

• 参数：
  • iterable：可迭代对象。
  • func：二元操作函数（默认求和）。
• 返回：累积结果的迭代器（类似前缀和）。
• 应用：区间和查询优化。

from itertools import accumulate
import operatorlist(accumulate([1, 2, 3, 4], operator.mul))  # 输出：[1, 2, 6, 24]（累乘）

islice(iterable, start, stop[, step])

• 参数：同列表切片 iterable[start:stop:step]，但支持迭代器。
• 返回：切片后的迭代器。
• 应用：分批处理大数据。

from itertools import islicelist(islice('ABCDEFG', 2, 6, 2))  # 输出：['C', 'E']

5. 实战参数技巧

动态生成排列组合

n = 3
r = 2
list(permutations(range(n), r))  # 生成 0~n-1 的 r 长度排列

批量处理数据

from itertools import islicedef batch_process(iterable, batch_size):it = iter(iterable)while batch := list(islice(it, batch_size)):yield batchdata = range(10)
for batch in batch_process(data, 3):print(batch)  # 分批输出：[0,1,2], [3,4,5], [6,7,8], [9]

二、应用场景 

1. 排列组合工具

场景：枚举所有可能性

import itertools# 全排列：[1,2,3] 的所有排列
permutations = list(itertools.permutations([1, 2, 3]))
# 输出: [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]# 组合：从4个元素中选2个的所有组合
combinations = list(itertools.combinations([1, 2, 3, 4], 2))
# 输出: [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]# 笛卡尔积：多个列表的所有组合
product = list(itertools.product(['a', 'b'], [1, 2]))
# 输出: [('a', 1), ('a', 2), ('b', 1), ('b', 2)]

应用：

• 枚举排列用于旅行商问题（TSP）的暴力解法。
• 组合常用于子集枚举（如背包问题的候选解）。

2. 无限迭代器

场景：周期性操作或生成序列

import itertools# 循环迭代器：周期性生成元素
cycle = itertools.cycle(['a', 'b', 'c'])
# 输出: a, b, c, a, b, c, ... (无限循环)# 计数器：生成递增序列
count = itertools.count(1, 2)  # 从1开始，步长为2
# 输出: 1, 3, 5, 7, ... (无限递增)

应用：

• 处理周期性输入（如游戏回合制模拟）。
• 生成测试用例的无限数据流。

3. 迭代器拼接与分组

场景：数据预处理

import itertools# 拼接多个迭代器
chained = itertools.chain([1, 2], [3, 4])
# 输出: 1, 2, 3, 4# 分组（需先排序）
data = [('A', 1), ('A', 2), ('B', 3)]
groups = itertools.groupby(data, key=lambda x: x[0])
# 输出: {'A': [(1, 2)], 'B': [(3,)]}

应用：

• 合并多个输入源（如多路归并排序）。
• 分组统计（如按类别聚合数据）。

4. 组合生成器

场景：高效生成特定模式

import itertools# 压缩：过滤掉对应掩码为False的元素
compressed = list(itertools.compress('ABCDEF', [1, 0, 1, 0, 1, 1]))
# 输出: ['A', 'C', 'E', 'F']# 累积操作（类似前缀和）
accumulated = list(itertools.accumulate([1, 2, 3, 4], lambda x, y: x + y))
# 输出: [1, 3, 6, 10]

应用：

• 快速过滤无效数据（如筛选合法坐标）。
• 前缀和优化区间查询问题。

5. 算法竞赛实战案例

1.计算所有子数组的和（前缀和优化）

问题描述：
给定数组 arr = [1, 2, 3, 4]，计算任意子数组 arr[i:j+1] 的和（例如 i=1, j=3 时，子数组为 [2, 3, 4]，和为 9）。

核心思路：前缀和数组

使用 itertools.accumulate 生成前缀和数组，将子数组和查询优化到 O(1) 时间复杂度。

步骤解析：

1. 构造前缀和数组：
   原数组 arr 前添加一个 0，再计算累积和。

   arr = [1, 2, 3, 4]
   prefix_sum = list(accumulate([0] + arr))  # [0, 1, 3, 6, 10]
   

    
   • prefix_sum[k] 表示原数组前 k 个元素的和（即 arr[0] + arr[1] + ... + arr[k-1]）。

2. 查询子数组和：
   子数组 arr[i:j+1] 的和可表示为：

   sum_i_j = prefix_sum[j+1] - prefix_sum[i]
   

    
   • 例如，查询 i=1, j=3（子数组 [2, 3, 4]）：
     prefix_sum[4] - prefix_sum[1] = 10 - 1 = 9。

完整代码

from itertools import accumulatearr = [1, 2, 3, 4]
prefix_sum = list(accumulate([0] + arr))  # [0, 1, 3, 6, 10]# 查询子数组 arr[1:3+1] 的和
i, j = 1, 3
sum_i_j = prefix_sum[j+1] - prefix_sum[i]
print(f"子数组和: {sum_i_j}")  # 输出: 9

2.生成格雷码（Gray Code）

问题描述：
格雷码是二进制数字系统，相邻两个数值仅有一位不同。生成 n 位格雷码序列（例如 n=3 时，序列为 [0, 1, 3, 2, 6, 7, 5, 4]）。

核心思路：位运算公式

使用公式 G(n) = n ^ (n >> 1) 直接生成格雷码，其中 ^ 是异或运算，>> 是右移运算。

步骤解析：

1. 公式推导：
   对于整数 n，其格雷码为 n 与 n 右移一位的异或结果。

   gray_code = n ^ (n >> 1)
   

    
   • 例如，n=3（二进制 11）：
     3 ^ (3 >> 1) = 11₂ ^ 01₂ = 10₂ = 2，因此格雷码为 2。

2. 生成序列：
   遍历 0 到 2^n - 1，对每个数应用公式：

   gray_codes = [a ^ (a >> 1) for a in range(2**n)]
   

完整代码

from itertools import productn = 3
gray_codes = [a ^ (a >> 1) for a in range(2**n)]
print(f"{n}位格雷码序列: {gray_codes}")
# 输出: [0, 1, 3, 2, 6, 7, 5, 4]

格雷码（Gray Code）是一种二进制编码方式，其特点是相邻的两个码字之间只有一位不同。