itertools —— Python 最强迭代器工具箱实战

1. 知识点简介

itertools 是 Python 标准库中最被低估的模块之一。它提供了一组高效、内存友好 的迭代器工具，让你用声明式的组合替代手写循环。核心设计哲学：延迟计算、流式处理、组合复用。

一句话总结：itertools 在手，循环代码再也不用从头搓。

---

2. 三大基石：无限迭代器

2.1 count(start=0, step=1) ------ 等差数列

from itertools import count

# 从 10 开始，步长 2
for i in count(10, 2):
    if i > 20:
        break
    print(i, end=" ")  # 10 12 14 16 18 20

实战：搭配 zip 做自动编号

from itertools import count

fruits = ["apple", "banana", "cherry"]
for idx, fruit in zip(count(1), fruits):
    print(f"{idx}. {fruit}")
# 1. apple
# 2. banana
# 3. cherry

比 enumerate(fruits, start=1) 的好处是：count 是迭代器，适合更复杂的 chaining 场景。

2.2 cycle(iterable) ------ 无限循环

from itertools import cycle
import time

status_indicators = cycle(["⏳", "🔄", "⚙️", "📦"])
for i, icon in enumerate(cycle(status_indicators)):
    if i > 20:
        break
    print(f"\r{icon} 处理中...", end="")
    time.sleep(0.2)

实战：轮询分配任务

from itertools import cycle

backends = cycle(["server-a", "server-b", "server-c"])

def get_next_backend():
    return next(backends)

for _ in range(5):
    print(get_next_backend())  # server-a, server-b, server-c, server-a, server-b

2.3 repeat(object, times=None) ------ 重复

from itertools import repeat

# 重复 5 次
for _ in repeat("doing work", 5):
    print("⚡", end=" ")  # ⚡ ⚡ ⚡ ⚡ ⚡

# 搭配 map 初始化
list(map(pow, range(5), repeat(3)))  # [0, 1, 8, 27, 64] ← 0³, 1³, 2³, 3³, 4³

---

3. 数据处理五虎将

3.1 chain(*iterables) ------ 扁平拼接多个可迭代对象

from itertools import chain

# 代替 a + b（不用创建中间列表）
a = [1, 2, 3]
b = [4, 5, 6]
c = ["x", "y"]

for item in chain(a, b, c):
    print(item, end=" ")  # 1 2 3 4 5 6 x y

# 展平嵌套列表（一层）
nested = [[1, 2], [3, 4], [5]]
flattened = list(chain.from_iterable(nested))  # [1, 2, 3, 4, 5]

chain vs +：+ 会创建新列表，chain 是迭代器，O(1) 空间。

3.2 compress(data, selectors) ------ 按布尔掩码过滤

from itertools import compress

data = ["正常", "异常", "正常", "严重", "正常"]
mask = [True, False, True, True, False]

filtered = list(compress(data, mask))  # ['正常', '正常', '严重']

# 实战：按条件筛选
scores = [88, 45, 92, 67, 73]
passing = [s >= 60 for s in scores]
passing_scores = list(compress(scores, passing))  # [88, 92, 67, 73]

3.3 dropwhile / takewhile ------ 条件式截取

from itertools import dropwhile, takewhile

data = [1, 2, 3, 4, 5, 1, 2, 3]

# 跳过开头 <= 3 的元素，遇到 > 3 才开始保留
list(dropwhile(lambda x: x <= 3, data))  # [4, 5, 1, 2, 3]

# 保留开头 <= 3 的元素，遇到 > 3 就停止
list(takewhile(lambda x: x <= 3, data))  # [1, 2, 3]

实战：跳过日志文件头部注释

# 假设日志内容
lines = ["# 开始时间: 2026-06-18", "# 版本: 2.0", "", "data1", "data2"]

from itertools import dropwhile
data_lines = list(dropwhile(lambda l: l.startswith("#") or l == "", lines))
# ['data1', 'data2']

3.4 filterfalse(predicate, iterable) ------ 反过滤

from itertools import filterfalse

nums = [1, 2, 3, 4, 5, 6]

# 选出「不是偶数」的数 → 奇数
list(filterfalse(lambda x: x % 2 == 0, nums))  # [1, 3, 5]
# 等价于
list(filter(lambda x: x % 2 != 0, nums))       # [1, 3, 5]

# ⚡ 返回不满足条件的元素，避免写双重否定

3.5 accumulate(iterable, func=operator.add) ------ 累加 / 前缀运算

from itertools import accumulate
import operator

# 前缀和
list(accumulate([1, 2, 3, 4, 5]))  # [1, 3, 6, 10, 15]

# 前缀积
list(accumulate([1, 2, 3, 4, 5], operator.mul))  # [1, 2, 6, 24, 120]

# 自定义累积：计算最大值变化
list(accumulate([5, 3, 8, 2, 9], max))  # [5, 5, 8, 8, 9]

# 实战：计算 running total 并跳过异常值
prices = [100, 200, -999, 300, 150]

def safe_add(acc, x):
    return acc + x if x > 0 else acc

list(accumulate(prices, safe_add))  # [100, 300, 300, 600, 750]

---

4. 排列组合三剑客

from itertools import product, permutations, combinations, combinations_with_replacement

letters = ["A", "B", "C"]
nums = [1, 2]

# 笛卡尔积（嵌套循环 flat 版）
list(product(letters, nums))
# [('A',1), ('A',2), ('B',1), ('B',2), ('C',1), ('C',2)]

# 排列（有序，不重复取）
list(permutations(letters, 2))
# [('A','B'), ('A','C'), ('B','A'), ('B','C'), ('C','A'), ('C','B')]

# 组合（无序，不重复取）
list(combinations(letters, 2))
# [('A','B'), ('A','C'), ('B','C')]

# 组合（允许重复取）
list(combinations_with_replacement([1, 2], 2))
# [(1, 1), (1, 2), (2, 2)]

实战：测试用例组合生成

from itertools import product

test_params = {
    "user_role": ["admin", "user", "guest"],
    "is_login": [True, False],
    "page": ["home", "profile", "settings"],
}

# 生成 3×2×3 = 18 个测试组合
for combo in product(*test_params.values()):
    params = dict(zip(test_params.keys(), combo))
    print(f"测试: {params}")
    # run_test(**params)

---

5. 分组神器：groupby

from itertools import groupby

data = [
    ("fruit", "apple"), ("fruit", "banana"),
    ("drink", "coffee"), ("drink", "tea"),
    ("fruit", "cherry"),  # ⚠️ 注意：groupby 要求已排序！
]

# 先排序再分组
data.sort(key=lambda x: x[0])
for category, items in groupby(data, key=lambda x: x[0]):
    print(f"{category}: {[item[1] for item in items]}")
# drink: ['coffee', 'tea']
# fruit: ['apple', 'banana', 'cherry']

⚠️ 核心坑 ：groupby 只对连续相同的元素分组，不是 SQL 的 GROUP BY。务必预排序。

# ❌ 未排序的后果
data = [("a", 1), ("b", 1), ("a", 2)]
for k, g in groupby(data, key=lambda x: x[0]):
    print(k, list(g))
# a [('a', 1)]   ← 第一个 a 组
# b [('b', 1)]
# a [('a', 2)]   ← 第二个 a 组！同一个 key 分成了两组

---

6. 组合技：tee、slice、starmap

from itertools import tee, islice, starmap

# tee：一个迭代器拆成多个独立副本
it = (x for x in range(5))
it1, it2 = tee(it, 2)
list(it1)  # [0, 1, 2, 3, 4]
list(it2)  # [0, 1, 2, 3, 4]  ← 互不影响

# islice：迭代器切片（不创建中间列表）
list(islice(range(100), 10, 20, 2))  # [10, 12, 14, 16, 18]
# 对比 list(range(100))[10:20:2] ← 创建了 100 个元素的临时列表

# starmap：带 * 展开的 map
data = [(2, 3), (4, 5), (6, 7)]
list(starmap(pow, data))  # [8, 1024, 279936]  ← pow(2,3), pow(4,5), pow(6,7)

---

7. 避坑指南

坑点	说明	解决
❌ groupby 未排序	返回的不是全量分组，而是「连续段」分组	分组前先 sorted(data, key=...)
❌ 迭代器已耗尽	itertools 返回的都是迭代器，只能遍历一次	需要重复用 → list() 转存或 tee()
❌ 无限迭代器没有终止条件	count() / cycle() 没 break 会死循环	搭配 takewhile / islice 限流
❌ 错误使用 tee	tee 的副本会共享内存，如果一个消耗快一个消耗慢，数据会在内存积压	副本之间保持同步消费，或用 list()
❌ 忽视性能优势	chain 写成 a + b、compress 写成列表推导	优先用 itertools 节省内存

---

8. 总结

• itertools 不是炫技工具，是解决常见迭代问题的标准答案
• 记住这几个高频组合：chain（展平）、product（笛卡尔积）、groupby（分组）、accumulate（前缀计算）、compress（掩码过滤）
• 所有工具都是惰性迭代器 ，搭配 islice / takewhile 可以无限流式处理
• 编写数据处理管道时，尽量用 itertools 组合替代手写循环，代码更短、含义更清晰、内存更省