字数 1212,阅读大约需 7 分钟
Python itertools 全解析

你写循环的时候,是不是还在老老实实地 for i in range(len(lst))?遇到嵌套循环就疯狂缩进,看到去重就 set() 一把梭?朋友,Python 标准库里有个宝藏模块叫 itertools,专门治你这种"只会 for 的人"。
这篇文章把 itertools 里最实用的函数一网打尽,看完你写循环的姿势至少优雅三个档次。
无限迭代器:count、cycle、repeat
这三个函数生成无限序列,配合 islice 食用效果最佳。
count(start, step) 从 start 开始按步长无限递增,比 while 循环优雅得多:
from itertools import count, islice
# 取前5个偶数
evens = list(islice(count(0, 2), 5))
# [0, 2, 4, 6, 8]
cycle(iterable) 无限循环一个序列,做轮询调度特别方便:
from itertools import cycle
servers = ['A', 'B', 'C']
round_robin = cycle(servers)
# next(round_robin) → A, B, C, A, B, C ...
repeat(obj, times) 重复一个元素,比 [obj] * n 更省内存:
from itertools import repeat
list(repeat('hello', 3))
# ['hello', 'hello', 'hello']
累积操作:accumulate
还在手写累加?accumulate 默认做累加,但你也可以传自定义函数:
from itertools import accumulate
import operator
nums = [1, 2, 3, 4, 5]
# 累加
list(accumulate(nums))
# [1, 3, 6, 10, 15]
# 累乘
list(accumulate(nums, operator.mul))
# [1, 2, 6, 24, 120]
# 最大值前缀
list(accumulate(nums, max))
# [1, 2, 3, 4, 5]
算前缀和、前缀积、前缀最值,一行搞定。你之前写的那些循环,删了吧。
链接与选择:chain、compress、filterfalse
chain(*iterables) 把多个可迭代对象串成一条链,比 sum(lists, []) 快且省内存:
from itertools import chain
list1 = [1, 2]
list2 = [3, 4]
list3 = [5, 6]
list(chain(list1, list2, list3))
# [1, 2, 3, 4, 5, 6]
# chain.from_iterable 处理嵌套列表更优雅
nested = [[1, 2], [3, 4], [5, 6]]
list(chain.from_iterable(nested))
# [1, 2, 3, 4, 5, 6]
compress(data, selectors) 按布尔序列筛选数据,比列表推导式更直观:
from itertools import compress
data = ['a', 'b', 'c', 'd']
mask = [True, False, True, False]
list(compress(data, mask))
# ['a', 'c']
filterfalse(predicate, iterable) 是 filter 的反义词,保留不满足条件的元素:
from itertools import filterfalse
nums = [1, 2, 3, 4, 5, 6]
list(filterfalse(lambda x: x % 2 == 0, nums))
# [1, 3, 5]
排列组合:product、permutations、combinations
这三个函数是排列组合的终极武器,比手写回溯简洁一百倍。
product(*iterables, repeat=1) 计算笛卡尔积,嵌套循环的克星:
from itertools import product
colors = ['红', '蓝']
sizes = ['S', 'M', 'L']
list(product(colors, sizes))
# [('红', 'S'), ('红', 'M'), ('红', 'L'), ('蓝', 'S'), ('蓝', 'M'), ('蓝', 'L')]
# 骰子所有可能
list(product(range(1, 7), repeat=2))
# 36种组合
permutations(iterable, r) 全排列:
from itertools import permutations
list(permutations('ABC', 2))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
combinations(iterable, r) 组合(不重复,不分顺序):
from itertools import combinations
list(combinations(range(1, 5), 2))
# [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
combinations_with_replacement 允许重复选择:
from itertools import combinations_with_replacement
list(combinations_with_replacement('ABC', 2))
# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]
分组与切片:groupby、islice、tee
groupby(iterable, key) 按键函数分组,注意必须先排序:
from itertools import groupby
students = [
{'name': '张三', 'grade': 'A'},
{'name': '李四', 'grade': 'B'},
{'name': '王五', 'grade': 'A'},
{'name': '赵六', 'grade': 'C'},
{'name': '钱七', 'grade': 'B'},
]
# 必须先按分组键排序!
students.sort(key=lambda x: x['grade'])
for grade, group in groupby(students, key=lambda x: x['grade']):
names = [s['name'] for s in group]
print(f'{grade}: {names}')
# A: ['张三', '王五']
# B: ['李四', '钱七']
# C: ['赵六']
islice(iterable, start, stop, step) 对迭代器切片,不用转成列表:
from itertools import islice
# 读取大文件前10行
with open('big_file.txt') as f:
first_10 = list(islice(f, 10))
tee(iterable, n) 从一个迭代器复制出 n 个独立迭代器:
from itertools import tee
it = iter(range(5))
it1, it2 = tee(it, 2)
list(it1) # [0, 1, 2, 3, 4]
list(it2) # [0, 1, 2, 3, 4]
实战:用 itertools 重构你的代码
来个真实场景。假设你要生成一个配置文件的所有参数组合,还要过滤掉不合法的组合:
from itertools import product, filterfalse
hosts = ['localhost', '192.168.1.1']
ports = [80, 443, 8080]
protocols = ['http', 'https']
# 所有组合
configs = product(hosts, ports, protocols)
# 过滤不合法:http 不能用 443,https 不该用 80
invalid = lambda c: (c[2] == 'http' and c[1] == 443) or (c[2] == 'https' and c[1] == 80)
valid_configs = list(filterfalse(invalid, configs))
# 8个组合 → 过滤后6个
如果用传统嵌套循环?三层 for + if,缩进能把你逼疯。用 itertools,一行生成,一行过滤,清爽。
性能小贴士
itertools 最大的优势不是花哨,而是惰性求值——它返回迭代器,只在需要时才计算下一个值。处理大数据集时,这意味着零额外内存开销。
几个原则:
- • 能用
chain 就别 + 拼列表,前者不创建中间对象 - • 能用
islice 就别 list(it)[0:10],前者不会把整个迭代器跑完 - • 能用
takewhile/dropwhile 就别在循环里 break,语义更清晰
标准库不是摆设,itertools 更是 Python 里被低估得最狠的模块之一。下次写循环之前,先翻翻这个模块,大概率你想要的它已经有了。