Python中的迭代器工具库：itertools模块详解

1. 模块概述

itertools 是 Python 标准库中的一个模块，专门用于创建和操作迭代器的工具集合。它提供了一系列高效、内存友好的函数，能够组合出复杂的迭代逻辑，常用于处理大数据流、组合数学问题、序列操作等场景。这些函数大多返回迭代器对象，采用惰性求值（lazy evaluation）方式，不会一次性生成所有结果，从而节省内存。

2. 核心函数分类与详解

itertools 中的函数可分为三大类：无限迭代器、有限迭代器和组合迭代器。

2.1 无限迭代器

这些迭代器会无限产生值，必须通过外部条件（如 break 或切片）来终止。

• count(start=0, step=1)
  从 start 开始，以 step 为步长无限递增。
  示例：

  import itertools
  for i in itertools.count(10, 2):
      if i > 20: break
      print(i)  # 输出：10, 12, 14, 16, 18, 20
  

• cycle(iterable)
  无限重复给定可迭代对象中的元素。
  示例：

  c = itertools.cycle('AB')
  for _ in range(4):
      print(next(c))  # 输出：A, B, A, B
  

• repeat(object, times=None)
  重复生成 object。如果指定 times，则重复固定次数；否则无限重复。
  示例：

  list(itertools.repeat('X', 3))  # 输出：['X', 'X', 'X']
  

2.2 有限迭代器

这些迭代器在输入耗尽时停止。

• accumulate(iterable, func=operator.add)
  对输入序列进行累积计算（类似前缀和）。默认使用加法，可自定义二元函数。
  示例：

  list(itertools.accumulate([1, 2, 3, 4]))  # 输出：[1, 3, 6, 10]
  import operator
  list(itertools.accumulate([2, 3, 4], operator.mul))  # 输出：[2, 6, 24]
  

• chain(*iterables)
  将多个可迭代对象连接成一个迭代器，按顺序遍历。
  示例：

  list(itertools.chain('AB', [1, 2]))  # 输出：['A', 'B', 1, 2]
  

• chain.from_iterable(iterable)
  将可迭代对象中的每个元素（本身也是可迭代的）展开连接。
  示例：

  list(itertools.chain.from_iterable(['AB', 'CD']))  # 输出：['A', 'B', 'C', 'D']
  

• compress(data, selectors)
  根据 selectors（布尔序列）筛选 data 中的元素。
  示例：

  list(itertools.compress('ABCD', [1, 0, 1, 0]))  # 输出：['A', 'C']
  

• dropwhile(predicate, iterable)
  跳过满足 predicate 的元素，直到遇到第一个不满足的，然后返回剩余所有元素。
  示例：

  list(itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 2]))  # 输出：[6, 2]
  

• takewhile(predicate, iterable)
  返回满足 predicate 的元素，直到遇到第一个不满足的为止。
  示例：

  list(itertools.takewhile(lambda x: x < 5, [1, 4, 6, 2]))  # 输出：[1, 4]
  

• filterfalse(predicate, iterable)
  返回不满足 predicate 的所有元素（与内置 filter 相反）。
  示例：

  list(itertools.filterfalse(lambda x: x % 2, [1, 2, 3, 4]))  # 输出：[2, 4]
  

• groupby(iterable, key=None)
  将连续且具有相同 key 值的元素分组。注意：需先按 key 排序才能正确分组。
  示例：

  data = sorted([('a', 1), ('b', 2), ('a', 3)], key=lambda x: x[0])
  for key, group in itertools.groupby(data, key=lambda x: x[0]):
      print(key, list(group))
  # 输出：
  # a [('a', 1), ('a', 3)]
  # b [('b', 2)]
  

• islice(iterable, start, stop, step=1)
  对迭代器进行切片，类似列表切片，但支持无限迭代器。
  示例：

  list(itertools.islice('ABCDEFG', 2, None, 2))  # 输出：['C', 'E', 'G']
  

• starmap(function, iterable)
  将可迭代对象中的每个元素（本身是元组）解包后传给 function。
  示例：

  list(itertools.starmap(pow, [(2, 3), (3, 2)]))  # 输出：[8, 9]
  

• tee(iterable, n=2)
  将一个迭代器拆分成 n 个独立的迭代器。注意：原迭代器被消耗后不应再使用。
  示例：

  a, b = itertools.tee(range(3), 2)
  list(a)  # 输出：[0, 1, 2]
  list(b)  # 输出：[0, 1, 2]
  

• zip_longest(*iterables, fillvalue=None)
  类似内置 zip，但以最长的可迭代对象为准，不足部分用 fillvalue 填充。
  示例：

  list(itertools.zip_longest('AB', [1, 2, 3], fillvalue='-'))
  # 输出：[('A', 1), ('B', 2), ('-', 3)]
  

2.3 组合迭代器

用于组合数学中的排列、组合等操作。

• product(*iterables, repeat=1)
  计算笛卡尔积（相当于多层嵌套循环）。
  示例：

  list(itertools.product('AB', repeat=2))  # 输出：[('A', 'A'), ('A', 'B'), ('B', 'A'), ('B', 'B')]
  

• permutations(iterable, r=None)
  返回长度为 r 的所有排列，顺序敏感且不含重复元素（若 r 未指定，则等于可迭代对象长度）。
  示例：

  list(itertools.permutations('AB', 2))  # 输出：[('A', 'B'), ('B', 'A')]
  

• combinations(iterable, r)
  返回长度为 r 的所有组合，顺序不敏感且不含重复元素。
  示例：

  list(itertools.combinations('ABC', 2))  # 输出：[('A', 'B'), ('A', 'C'), ('B', 'C')]
  

• combinations_with_replacement(iterable, r)
  返回长度为 r 的所有组合，允许元素重复。
  示例：

  list(itertools.combinations_with_replacement('AB', 2))
  # 输出：[('A', 'A'), ('A', 'B'), ('B', 'B')]
  

3. 核心特点与最佳实践

1. 惰性求值：所有函数返回迭代器，避免一次性加载全部数据，适合处理大规模数据流。
2. 内存高效：迭代器逐个生成元素，内存占用恒定。
3. 可组合性：多个 itertools 函数可链式调用，构建复杂数据处理管道。
4. 性能优势：底层用 C 实现，比纯 Python 循环更快。

使用示例：生成所有三位数的排列（数字不重复）：

import itertools
perms = itertools.permutations(range(10), 3)
for p in perms:
    num = p[0]*100 + p[1]*10 + p[2]
    if num > 200:
        print(num)
        break

4. 常见应用场景

• 数据流处理：连接、过滤、转换大型数据集。
• 组合问题：密码破解、游戏状态枚举、实验设计。
• 窗口计算：与 islice 或 tee 结合实现滑动窗口。
• 分组聚合：对排序后的数据按关键字段分组（需先排序）。

5. 注意事项

• 无限迭代器需有终止条件，否则会导致无限循环。
• groupby 需先排序，因为它只对连续相同键值分组。
• tee 的迭代器共享数据，原迭代器被消耗后，复制出的迭代器也会受影响。

通过掌握 itertools，你可以用更简洁、高效的方式处理迭代任务，减少显式循环和中间列表的创建，提升代码性能和可读性。