Python 中 itertools 模块的常用工具
python 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
**实战场景**:给数据流添加自动递增的索引。
```python
data = ['apple', 'banana', 'cherry']
for index, item in zip(count(1), data):
print(f"{index}. {item}")
# 1. apple
# 2. banana
# 3. cherry
2. cycle(iterable)
cycle 会无限循环地重复输出传入的可迭代对象中的元素。
from itertools import cycle
colors = cycle(['red', 'green', 'blue'])
for i in range(5):
print(next(colors), end=' ') # 输出:red green blue red green
实战场景:轮询分发任务到多个服务器,或生成循环的颜色标记。
3. repeat(object, times=None)
repeat 重复返回同一个对象。若指定 times 参数,则重复指定次数;否则无限重复。
from itertools import repeat
# 无限重复 10
# for x in repeat(10): ...
# 重复 3 次 'Hello'
for msg in repeat('Hello', 3):
print(msg, end=' ') # 输出:Hello Hello Hello
与 map 联用:快速生成固定值列表。
list(map(pow, range(5), repeat(2)))
# 计算 0^2, 1^2, 2^2, 3^2, 4^2 => [0, 1, 4, 9, 16]
二、终止于最短输入序列的迭代器:数据处理的瑞士军刀
这是 itertools 的核心功能。它们会依据最短输入序列的结束来终止迭代,防止越界。
1. accumulate(iterable, func=operator.add)
accumulate 返回累积结果序列。默认执行累加求和,可自定义函数。
from itertools import accumulate
import operator
nums = [1, 2, 3, 4, 5]
# 默认累加
print(list(accumulate(nums))) # [1, 3, 6, 10, 15]
# 累乘
print(list(accumulate(nums, operator.mul))) # [1, 2, 6, 24, 120]
# 累积最大值
print(list(accumulate(nums, max))) # [1, 2, 3, 4, 5]
2. chain(*iterables)
chain 将多个可迭代对象无缝拼接成一个序列。
from itertools import chain
letters = ['a', 'b']
numbers = [1, 2]
flags = (True, False)
combined = chain(letters, numbers, flags)
print(list(combined)) # ['a', 'b', 1, 2, True, False]
chain.from_iterable:当需要拼接的对象是一个嵌套可迭代对象时使用。
nested = [['A', 'B'], ['C', 'D']]
flat = chain.from_iterable(nested)
print(list(flat)) # ['A', 'B', 'C', 'D']
3. compress(data, selectors)
compress 根据 selectors 中元素的真假值,过滤 data 中对应的元素。
from itertools import compress
words = ['Python', 'is', 'awesome', '!']
selector = [1, 0, 1, 0] # 0 和 1 可直接作为布尔值使用
print(list(compress(words, selector))) # ['Python', 'awesome']
4. dropwhile 与 takewhile
dropwhile(predicate, iterable):从第一个令predicate为False的元素开始返回所有后续元素。takewhile(predicate, iterable):返回满足条件的前缀元素,直到碰到第一个不满足的元素。
from itertools import dropwhile, takewhile
scores = [55, 60, 70, 80, 45, 90]
# 跳过开头的低分,保留从第一个及格分数开始的所有记录
passed = dropwhile(lambda x: x < 60, scores)
print(list(passed)) # [60, 70, 80, 45, 90]
# 只取开头的连续及格分数
front_passed = takewhile(lambda x: x >= 60, scores)
print(list(front_passed)) # []
5. filterfalse(predicate, iterable)
返回使 predicate 为 False 的所有元素,与内置 filter 互补。
from itertools import filterfalse
nums = range(10)
# 保留所有偶数(即 filter 的结果)
evens = list(filter(lambda x: x%2==0, nums)) # [0, 2, 4, 6, 8]
# 保留所有奇数
odds = list(filterfalse(lambda x: x%2==0, nums)) # [1, 3, 5, 7, 9]
6. groupby(iterable, key=None)
groupby 将连续相同的键值分组。注意:必须事先对数据排序,否则分组会不完整。
from itertools import groupby
pets = [('cat', 'Mimi'), ('dog', 'Rex'), ('cat', 'Tom'), ('dog', 'Spike')]
# 先按动物种类排序
pets_sorted = sorted(pets, key=lambda x: x[0])
for animal, group in groupby(pets_sorted, key=lambda x: x[0]):
names = [name for _, name in group]
print(f"{animal}: {', '.join(names)}")
# cat: Mimi, Tom
# dog: Rex, Spike
7. islice(iterable, start, stop, step)
如同列表切片,但适用于任意迭代器,且懒加载。参数与 slice 行为一致。
from itertools import islice
# 取斐波那契数列的第 5 到第 9 个元素(索引从 0 开始)
def fibonacci():
a, b = 0, 1
while True:
yield a
a, b = b, a + b
fib_slice = islice(fibonacci(), 5, 10)
print(list(fib_slice)) # [5, 8, 13, 21, 34]
8. pairwise(iterable) (Python 3.10+)
返回连续重叠的相邻元素对,相当于 (i0,i1),(i1,i2),...。
from itertools import pairwise
seq = 'ABCDE'
for a, b in pairwise(seq):
print(f"{a} -> {b}")
# A -> B
# B -> C
# C -> D
# D -> E
9. starmap(function, iterable)
当你的 iterable 每个元素本身是参数元组时,starmap 会将它们解包后传给函数。
from itertools import starmap
points = [(2, 5), (3, 1), (10, 4)]
results = starmap(pow, points) # 计算 2^5, 3^1, 10^4
print(list(results)) # [32, 3, 10000]
10. zip_longest(*iterables, fillvalue=None)
与 zip 不同,它会根据最长的输入序列迭代,缺失值用 fillvalue 填充。
from itertools import zip_longest
names = ['Alice', 'Bob']
ages = [25]
cities = ['NYC', 'LA', 'London']
combined = zip_longest(names, ages, cities, fillvalue='Unknown')
for item in combined:
print(item)
# ('Alice', 25, 'NYC')
# ('Bob', 'Unknown', 'LA')
# ('Unknown', 'Unknown', 'London')
三、组合生成器:轻松驾驭排列与组合
用于生成序列元素的排列、组合和笛卡尔积。
1. product(*iterables, repeat=1)
计算笛卡尔积,即所有可能的配对组合。repeat 参数允许一个迭代器自身乘几次。
from itertools import product
# 两个骰子的所有可能结果
dice = product(range(1, 3), range(1, 3))
print(list(dice))
# [(1, 1), (1, 2), (2, 1), (2, 2)]
# 三位二进制数(重复使用 '01' 三次)
bits = product('01', repeat=3)
print([''.join(b) for b in bits])
# ['000', '001', '010', '011', '100', '101', '110', '111']
2. permutations(iterable, r=None)
返回长度为 r 的所有可能排列(顺序相关)。若 r 未指定,则取全长。
from itertools import permutations
chars = 'AB'
# 长度为 2 的排列
print(list(permutations(chars))) # [('A', 'B'), ('B', 'A')]
# 长度为 1 的排列
print(list(permutations(chars, 1))) # [('A',), ('B',)]
3. combinations(iterable, r)
返回长度为 r 的所有组合,顺序无关且元素不重复。
from itertools import combinations
items = ['A', 'B', 'C']
print(list(combinations(items, 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'C')]
4. combinations_with_replacement(iterable, r)
允许单个元素在组合中重复出现。
from itertools import combinations_with_replacement
print(list(combinations_with_replacement('AB', 2)))
# [('A', 'A'), ('A', 'B'), ('B', 'B')]
实践提升:动手构建一个数据分析流水线
假设我们有一份用户行为日志,需要提取“连续三次失败登录”的时间窗口。我们可以结合 groupby 和 accumulate 快速实现。
import itertools
from datetime import datetime, timedelta
events = [
('user1', 'fail'), ('user1', 'fail'), ('user1', 'success'),
('user1', 'fail'), ('user1', 'fail'), ('user1', 'fail'),
('user2', 'fail'), ('user2', 'fail')
]
# 1. 按用户分组
for user, group in itertools.groupby(events, key=lambda x: x[0]):
actions = [a for _, a in group]
# 2. 使用 accumulate 构建连续失败计数的重置逻辑
# 当遇到 'success' 时计数归零,否则加 1
def fail_counter(acc, action):
return 0 if action == 'success' else acc + 1
fail_streaks = itertools.accumulate(actions, fail_counter)
# 3. 找出计数达到 3 的位置
for idx, streak in enumerate(fail_streaks):
if streak == 3:
print(f"{user} 在位置 {idx-2} 到 {idx} 连续三次登录失败,触发告警!")
break
输出:
user1 在位置 3 到 5 连续三次登录失败,触发告警!