Python itertools 迭代工具

学习itertools · 难度：高级 · +15XP

Python itertools 迭代工具

itertools 是 Python 标准库中提供高效迭代器构建块的模块。它包含一系列用于产生复杂迭代逻辑的函数，如组合、排列、笛卡尔积、无限迭代器等。所有函数返回的都是迭代器，内存效率极高。这些工具在算法竞赛、数据处理和函数式编程中非常实用。

无限迭代器

from itertools import count, cycle, repeat
# count(start, step) - 无限计数
for i in count(10, 2):
    if i > 20: break
    print(i, end=' ')  # 10 12 14 16 18 20
# cycle(iterable) - 无限循环
cnt = 0
for item in cycle(['A', 'B', 'C']):
    if cnt >= 6: break
    print(item, end=' ')  # A B C A B C
    cnt += 1
# repeat(obj, times) - 重复对象
print(list(repeat('Hello', 3)))  # ['Hello', 'Hello', 'Hello']

组合函数

from itertools import product, permutations, combinations, combinations_with_replacement
items = ['A', 'B', 'C']
# product - 笛卡尔积
print(list(product(items, repeat=2)))
# [('A','A'),('A','B'),('A','C'),('B','A'),...]
# permutations - 排列 (顺序有关)
print(list(permutations(items, 2)))
# [('A','B'),('A','C'),('B','A'),('B','C'),('C','A'),('C','B')]
# combinations - 组合 (顺序无关)
print(list(combinations(items, 2)))
# [('A','B'),('A','C'),('B','C')]
# combinations_with_replacement - 可重复组合
print(list(combinations_with_replacement(items, 2)))
# [('A','A'),('A','B'),('A','C'),('B','B'),('B','C'),('C','C')]

链式与过滤工具

from itertools import chain, compress, takewhile, dropwhile
# chain - 串联多个迭代器
print(list(chain([1, 2], ['a', 'b'], (True, False))))
# [1, 2, 'a', 'b', True, False]
# chain.from_iterable - 从嵌套结构展开
nested = [[1, 2], [3, 4], [5, 6]]
print(list(chain.from_iterable(nested)))  # [1, 2, 3, 4, 5, 6]
# compress(data, selectors) - 按条件过滤
data = ['a', 'b', 'c', 'd']
selectors = [True, False, True, False]
print(list(compress(data, selectors)))  # ['a', 'c']
# takewhile / dropwhile - 条件取/舍
nums = [1, 3, 5, 2, 4, 6]
print(list(takewhile(lambda x: x < 4, nums)))  # [1, 3]
print(list(dropwhile(lambda x: x < 4, nums)))  # [5, 2, 4, 6]

itertools 分类速查