Python itertools:循环的核武器

同事的三层嵌套把我看懵了

上周review代码，看到同事写了个三层嵌套循环找组合。我当时就血压上来了。😂 三个for套在一起，像俄罗斯套娃，我看着都头晕。他说"这样逻辑清晰"。清晰个锤子，这是2024年了兄弟。

itertools到底是个啥

itertools是标准库里的隐藏大佬。你写循环能遇到的问题，它基本都有解。combinations、permutations、chain、groupby... 这帮人凑一块儿，就是循环界的特种部队。不用安装，开箱即用。

import itertools

# 不用写嵌套循环了，一行搞定组合
items = ['A', 'B', 'C', 'D']
for combo in itertools.combinations(items, 2):
    print(combo)
# 输出: ('A','B') ('A','C') ('A','D') ('B','C') ...

combinations 救了我的命

那天我要从20个元素里选3个做测试。要是手写循环，我得写多少行？combinations(iterable, r)直接出所有r长度的组合。重点来了——它是惰性求值的，内存根本炸不了。

from itertools import combinations

# 从100个里选3个，C(100,3) = 161700种
# 但内存不会爆炸，因为返回的是迭代器
nums = range(100)
count = 0
for triplet in combinations(nums, 3):
if sum(triplet) % 7 == 0:
        count += 1
print(f"满足条件的有 {count} 个")
# 运行飞快，内存占用稳如老狗

permutations 让我少熬一个通宵

做全排列的时候，我之前傻傻地写递归。permutations一出来，递归代码可以直接删了。每个排列都是元组形式，长度就是你要的r。注意顺序不同算不同排列，跟组合不一样。

from itertools import permutations

# 4个字母全排列，4! = 24种
chars = '甲乙丙丁'
for p in permutations(chars, 3):
    print(''.join(p))

# 带过滤条件的排列筛选
valid = [p for p in permutations(range(1,5), 4)
if p[0] * p[1] > 6]
print(f"符合条件的排列数: {len(valid)}")

groupby 差点把我整不会了

groupby这个函数有个大坑：它只对相邻相同元素分组。我第一次用没排序，结果分组碎了一地，当场懵了。必须先sort再groupby，记住了。

from itertools import groupby

# 错误用法（不排序）
data = ['A', 'B', 'A', 'A', 'B']
for k, g in groupby(data):
    print(k, list(g))
# A ['A']  B ['B']  A ['A','A']  B ['B']  ← 碎了！

# 正确用法
sorted_data = sorted(data)
for k, g in groupby(sorted_data):
    print(k, len(list(g)))
# A 2  B 2  ← 这才是我要的

无限迭代器才是真的骚

count()、cycle()、repeat()这三个无限迭代器，我第一次见的时候在想：无限循环？这不作死吗。后来做数据填充、轮询、生成无限ID的时候，真香。😂

from itertools import count, cycle, repeat, islice

# count(start=10, step=2) → 10, 12, 14, 16...
for i in islice(count(10, 2), 5):
    print(i)  # 10 12 14 16 18

# cycle 无限循环列表
colors = cycle(['红', '绿', '蓝'])
for _ in range(6):
    print(next(colors))  # 红绿蓝红绿蓝

# repeat 重复同一个值
print(list(repeat('★', 5)))  # ['★', '★', '★', '★', '★']

chain 扁平化神器

你有好几个列表想一起遍历，别用+拼接，内存浪费。chain(*iterables)直接串起来，零拷贝。我做日志合并的时候全靠它。

from itertools import chain

log1 = ['ERROR: 磁盘满', 'WARN: 内存高']
log2 = ['INFO: 启动完成']
log3 = ['ERROR: 连接超时', 'ERROR: 超时']

# 比 log1+log2+log3 省内存
for entry in chain(log1, log2, log3):
if entry.startswith('ERROR'):
        print(entry)

# chain.from_iterable 更狠，展平嵌套列表
nested = [[1,2], [3,4], [5,6]]
print(list(chain.from_iterable(nested)))
# [1, 2, 3, 4, 5, 6]

我用这几个函数搞了个笛卡尔积

product()做笛卡尔积，替代多层嵌套循环的神器。我以前写SQL生成测试数据，两层循环套了80行。用product压缩到5行。☕

from itertools import product

# 笛卡尔积：每个属性组合都覆盖到
sizes = ['S', 'M', 'L']
colors = ['黑', '白', '灰']
styles = ['圆领', 'V领']

for item in product(sizes, colors, styles):
    print(f"SKU: {'-'.join(item)}")
# 3×3×2 = 18种组合，一行搞定