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 循环快。