跳到主要内容

itertools

问题

itertools 模块有哪些常用函数?如何用 itertools 高效处理迭代?

答案

无限迭代器

from itertools import count, cycle, repeat

# count:无限递增
for i in count(10, 2): # 10, 12, 14, 16, ...
if i > 20: break

# cycle:无限循环
colors = cycle(["red", "green", "blue"])
[next(colors) for _ in range(5)] # ['red', 'green', 'blue', 'red', 'green']

# repeat:重复
list(repeat("x", 3)) # ['x', 'x', 'x']

组合与排列

from itertools import product, permutations, combinations, combinations_with_replacement

# 笛卡尔积
list(product("AB", "12")) # [('A','1'), ('A','2'), ('B','1'), ('B','2')]

# 全排列
list(permutations("ABC", 2)) # [('A','B'), ('A','C'), ('B','A'), ...]

# 组合(不重复)
list(combinations("ABC", 2)) # [('A','B'), ('A','C'), ('B','C')]

# 组合(可重复)
list(combinations_with_replacement("AB", 2)) # [('A','A'), ('A','B'), ('B','B')]

常用工具

from itertools import chain, islice, groupby, accumulate, zip_longest

# chain:连接多个迭代器
list(chain([1, 2], [3, 4], [5])) # [1, 2, 3, 4, 5]

# islice:切片迭代器(不创建列表)
list(islice(range(100), 5, 10)) # [5, 6, 7, 8, 9]

# groupby:分组(需要先排序)
data = sorted([(1, "a"), (1, "b"), (2, "c")], key=lambda x: x[0])
for key, group in groupby(data, key=lambda x: x[0]):
print(key, list(group))
# 1 [(1, 'a'), (1, 'b')]
# 2 [(2, 'c')]

# accumulate:累积
list(accumulate([1, 2, 3, 4])) # [1, 3, 6, 10]

常见面试问题

Q1: 如何展平嵌套列表?

答案

from itertools import chain

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

Q2: itertools 的优势是什么?

答案

所有 itertools 函数返回迭代器(惰性求值),内存效率极高,适合处理大数据量。它们是 C 实现的,比纯 Python 循环快。

相关链接