itertools 模块

itertools 模块提供了一系列高性能的迭代器生成函数,这些函数可以:

  • 无缝组合(像拼积木一样)
  • 惰性计算(按需生成,不一次性占用内存)
  • 替代循环与列表推导式(更高效、更优雅)

导入方式:

1
import itertools

itertools 常见函数

类别 代表函数 功能说明
迭代器生成器 count(), cycle(), repeat() 创建无限序列
迭代器切片与选择 islice(), takewhile(), dropwhile(), filterfalse() 从迭代器中有条件地截取元素
迭代器组合 chain(), zip_longest(), product(), combinations() 组合或排列多个迭代器
分组与聚合 groupby(), accumulate() 对连续数据分组或累积
其他工具 tee(), starmap() 复制迭代器、展开元组调用函数

核心函数详解

itertools.chain():多个迭代器顺序连接

语法:

1
itertools.chain(*iterables)

功能: 按顺序把多个可迭代对象连接成一个连续序列。

示例:

1
2
3
4
5
6
7
from itertools import chain

a = [1, 2, 3]
b = ['x', 'y', 'z']

for item in chain(a, b):
print(item, end=' ')

输出:

1
1 2 3 x y z

变体:

1
itertools.chain.from_iterable(list_of_iterables)

当输入是一个“迭代器的迭代器”时使用。

1
2
3
data = [[1, 2], [3, 4], [5, 6]]
flat = chain.from_iterable(data)
print(list(flat)) # [1, 2, 3, 4, 5, 6]

应用场景:

  • 拼接多个序列
  • 将嵌套结构“打平”(flatten)
  • 构建数据管道

itertools.cycle():无限循环迭代

语法:

1
itertools.cycle(iterable)

功能: 重复循环遍历一个序列中的元素。

示例:

1
2
3
4
from itertools import cycle, islice

colors = cycle(['red', 'green', 'blue'])
print(list(islice(colors, 7))) # 取前7个

输出:

1
['red', 'green', 'blue', 'red', 'green', 'blue', 'red']

应用场景:

  • 实现循环任务(例如状态机、动画帧循环)
  • 无限模式生成(如颜色序列、轮询分配任务)

注意:cycle() 生成无限迭代器,请搭配 islice() 限制输出!

itertools.islice():对迭代器进行切片

语法:

1
itertools.islice(iterable, start, stop[, step])

功能:list[start:stop:step] 一样切片,但作用于任意可迭代对象,而且不会将其整体加载入内存。

示例:

1
2
3
4
5
from itertools import islice, count

nums = count(10) # 10, 11, 12, ...
result = islice(nums, 3, 8) # 从第3个开始(索引3),到索引8(不含)
print(list(result)) # [13, 14, 15, 16, 17]

应用场景:

  • 惰性读取大文件的前 N 行;
  • 跳过前几个元素;
  • 替代普通切片操作(适用于生成器)。

itertools.groupby():对连续相同元素分组

语法:

1
itertools.groupby(iterable, key=None)

功能: 将迭代器中相邻的相同(或相同 key 值)的元素分组。

注意:groupby 不会自动排序,它只对连续相同元素分组,若想按值分组,请先用 sorted()

示例 1:基础用法

1
2
3
4
5
from itertools import groupby

data = 'aaabbccccdd'
for key, group in groupby(data):
print(key, list(group))

输出:

1
2
3
4
a ['a', 'a', 'a']
b ['b', 'b']
c ['c', 'c', 'c', 'c']
d ['d', 'd']

示例 2:带 key 函数

1
2
3
4
5
6
7
8
9
10
11
people = [
{'name': 'Tom', 'age': 20},
{'name': 'Jerry', 'age': 20},
{'name': 'Anna', 'age': 25},
]

# 必须先按 key 排序
people.sort(key=lambda p: p['age'])

for age, group in groupby(people, key=lambda p: p['age']):
print(age, [p['name'] for p in group])

输出:

1
2
20 ['Tom', 'Jerry']
25 ['Anna']

应用场景:

  • 日志分组(按日期、分类等)
  • 数据聚合(类似 SQL 的 GROUP BY)
  • 统计连续相同元素(压缩编码)

更多实用函数

函数 功能 示例
count(start=0, step=1) 无限计数器 count(10, 2) → 10,12,14,…
repeat(obj, n=None) 重复生成同一对象 repeat("A", 3) → A,A,A
accumulate(iterable, func=operator.add) 累积运算 [1,2,3,4] → [1,3,6,10]
product(A, B) 笛卡尔积 A×B
permutations(A, r) 排列 (A,B)
combinations(A, r) 组合(无重复) (A,B)
filterfalse(func, iterable) 返回 func 为 False 的元素 过滤反逻辑
takewhile(func, iterable) 取条件为真时的元素(遇假停止)
dropwhile(func, iterable) 跳过条件为真时的元素(遇假开始)
tee(iterable, n=2) 复制 n 个独立迭代器

函数式“工具链”组合

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from itertools import chain, cycle, islice, groupby

# 示例:合并两个数据流,循环标记,并分组处理
a = [1, 2, 3]
b = [4, 5]
colors = cycle(['red', 'green']) # 无限循环

# 链接并配色
mixed = zip(chain(a, b), islice(colors, len(a) + len(b)))
print(list(mixed))
# [(1, 'red'), (2, 'green'), (3, 'red'), (4, 'green'), (5, 'red')]

# 对字符串进行分组
data = "aaabbccccddd"
groups = [(k, len(list(g))) for k, g in groupby(data)]
print(groups)
# [('a', 3), ('b', 2), ('c', 4), ('d', 3)]

这种组合思想 = “函数式流水线处理”。它能让你像 UNIX 管道一样操作数据流,保持简洁、高效、惰性。

itertools 设计哲学

特点 含义
惰性(Lazy evaluation) 不提前计算,节省内存
组合性(Composable) 各函数能自然衔接使用
可迭代性(Iterable everywhere) 统一接口,所有函数返回迭代器
函数式思维(Functional style) 无状态、无副作用

总结表

函数 类型 功能 常见用途
chain 组合 串联多个迭代器 多数据源合并
cycle 生成 无限循环迭代 循环任务分配
islice 选择 惰性切片 控制输出长度
groupby 聚合 相邻分组 数据统计、压缩
count 生成 无限计数 索引流生成
accumulate 聚合 累积计算 前缀和、累计乘积
product 组合 笛卡尔积 符号表、组合实验
tee 工具 拆分迭代器 多次消费同源流