Python 中 itertools 模块的常用工具

FreeGuideOnline 13阅读 2026-07-07

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. dropwhiletakewhile

  • dropwhile(predicate, iterable):从第一个令 predicateFalse 的元素开始返回所有后续元素。
  • 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)

返回使 predicateFalse 的所有元素,与内置 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')]

实践提升:动手构建一个数据分析流水线

假设我们有一份用户行为日志,需要提取“连续三次失败登录”的时间窗口。我们可以结合 groupbyaccumulate 快速实现。

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 连续三次登录失败,触发告警!