Python中的itertools模块常见函数用法示例

itertools ，迭代工具模块，提供了用于高效处理迭代器和组合问题的工具。

1. itertools.permutations(iterable, r=None)

• 功能：生成输入迭代器的所有可能排列。

• 参数：

  • iterable：输入的可迭代对象。
  • r：可选参数，指定排列的长度。如果不指定，则默认为输入迭代器的长度。

• 示例：

  from itertools import permutations
  
  # 生成所有排列
  for perm in permutations([1, 2, 3]):
      print(perm)
  # 输出：
  # (1, 2, 3)
  # (1, 3, 2)
  # (2, 1, 3)
  # (2, 3, 1)
  # (3, 1, 2)
  # (3, 2, 1)
  
  # 生成长度为2的排列
  for perm in permutations([1, 2, 3], 2):
      print(perm)
  # 输出：
  # (1, 2)
  # (1, 3)
  # (2, 1)
  # (2, 3)
  # (3, 1)
  # (3, 2)
  

• 应用场景：用于生成所有可能的排列组合，常用于排列问题、字符串排列等。

2. itertools.combinations(iterable, r)

• 功能：生成输入迭代器的所有可能组合。

• 参数：

  • iterable：输入的可迭代对象。
  • r：组合的长度。

• 示例：

  from itertools import combinations
  
  # 生成长度为2的组合
  for comb in combinations([1, 2, 3], 2):
      print(comb)
  # 输出：
  # (1, 2)
  # (1, 3)
  # (2, 3)
  

• 应用场景：用于生成所有可能的组合，常用于组合问题、子集问题等。

3. itertools.combinations_with_replacement(iterable, r)

• 功能：生成输入迭代器的所有可能组合，允许重复选择。

• 参数：

  • iterable：输入的可迭代对象。
  • r：组合的长度。

• 示例：

  from itertools import combinations_with_replacement
  
  # 生成长度为2的组合，允许重复选择
  for comb in combinations_with_replacement([1, 2, 3], 2):
      print(comb)
  # 输出：
  # (1, 1)
  # (1, 2)
  # (1, 3)
  # (2, 2)
  # (2, 3)
  # (3, 3)
  

• 应用场景：用于生成允许重复选择的组合，常用于组合问题中的重复选择场景。

4. itertools.product(*iterables, repeat=1)

• 功能：生成输入迭代器的笛卡尔积。

• 参数：

  • *iterables：多个可迭代对象。
  • repeat：可选参数，指定重复次数。

• 示例：

  from itertools import product
  
  # 生成两个列表的笛卡尔积
  for prod in product([1, 2], [3, 4]):
      print(prod)
  # 输出：
  # (1, 3)
  # (1, 4)
  # (2, 3)
  # (2, 4)
  
  # 生成一个列表的重复笛卡尔积
  for prod in product([1, 2], repeat=3):
      print(prod)
  # 输出：
  # (1, 1, 1)
  # (1, 1, 2)
  # (1, 2, 1)
  # (1, 2, 2)
  # (2, 1, 1)
  # (2, 1, 2)
  # (2, 2, 1)
  # (2, 2, 2)
  

• 应用场景：用于生成多个列表的笛卡尔积，常用于多维遍历、状态空间搜索等。

5. itertools.groupby(iterable, key=None)

• 功能：将相邻的重复元素挑出来放在一起。

• 参数：

  • iterable：输入的可迭代对象。
  • key：可选参数，用于指定分组的键函数。

• 示例：

  from itertools import groupby
  
  # 按值分组
  data = [1, 1, 2, 2, 3, 3]
  for key, group in groupby(data):
      print(key, list(group))
  # 输出：
  # 1 [1, 1]
  # 2 [2, 2]
  # 3 [3, 3]
  
  # 按键函数分组
  data = [("a", 1), ("a", 2), ("b", 3), ("b", 4)]
  for key, group in groupby(data, key=lambda x: x[0]):
      print(key, list(group))
  # 输出：
  # a [('a', 1), ('a', 2)]
  # b [('b', 3), ('b', 4)]
  

• 应用场景：用于对数据进行分组，常用于数据处理、统计等。

6. itertools.accumulate(iterable, func=None, *, initial=None)

• 功能：生成输入迭代器的累积和或累积结果。

• 参数：

  • iterable：输入的可迭代对象。
  • func：可选参数，指定累积操作的函数，默认为加法。
  • initial：可选参数，指定初始值。

• 示例：

  from itertools import accumulate
  
  # 累加
  data = [1, 2, 3, 4]
  print(list(accumulate(data)))  # 输出：[1, 3, 6, 10]
  
  # 累乘
  from operator import mul
  print(list(accumulate(data, mul)))  # 输出：[1, 2, 6, 24]
  
  # 带初始值
  print(list(accumulate(data, initial=10)))  # 输出：[10, 11, 13, 16, 20]
  

• 应用场景：用于生成累积和或累积结果，常用于动态规划、前缀和等。

7. itertools.chain(*iterables)

• 功能：将多个迭代器串联为一个长迭代器。

• 参数：

  • *iterables：多个可迭代对象。

• 示例：

  from itertools import chain
  
  # 串联多个列表
  data1 = [1, 2, 3]
  data2 = [4, 5, 6]
  print(list(chain(data1, data2)))  # 输出：[1, 2, 3, 4, 5, 6]
  

• 应用场景：用于将多个迭代器合并为一个，常用于数据拼接。

8. itertools.zip_longest(*iterables, fillvalue=None)

• 功能：将多个迭代器的元素配对，长度不足时用指定值填充。

• 参数：

  • *iterables：多个可迭代对象。
  • fillvalue：可选参数，指定填充值。

• 示例：

  from itertools import zip_longest
  
  # 配对并填充
  data1 = [1, 2, 3]
  data2 = [4, 5]
  print(list(zip_longest(data1, data2, fillvalue=0)))  # 输出：[(1, 4), (2, 5), (3, 0)]
  

• 应用场景：用于处理长度不一致的迭代器，常用于数据对齐。

总结

• 排列和组合问题：使用 permutations、combinations 和 combinations_with_replacement。
• 多维遍历：使用 product。
• 数据分组：使用 groupby。
• 累积计算：使用 accumulate。
• 数据拼接：使用 chain。
• 数据对齐：使用 zip_longest。