Python 内置模块 itertools

无限迭代器

count() - 无限计数器

从指定数字开始无限计数，可设置步长。

import itertools

# 从0开始计数
counter = itertools.count()
for i, num in enumerate(counter):
    if i >= 5:
        break
    print(num)  # 输出: 0, 1, 2, 3, 4

# 从10开始，步长为2
counter2 = itertools.count(10, 2)
result = [next(counter2) for _ in range(4)]
print(result)  # [10, 12, 14, 16]

cycle() - 循环迭代

无限循环一个可迭代对象的元素。

import itertools

colors = ['red', 'green', 'blue']
color_cycle = itertools.cycle(colors)

result = [next(color_cycle) for _ in range(7)]
print(result)  # ['red', 'green', 'blue', 'red', 'green', 'blue', 'red']

repeat() - 重复元素

重复一个值指定次数或无限重复。

import itertools

# 重复5次
numbers = list(itertools.repeat(42, 5))
print(numbers)  # [42, 42, 42, 42, 42]

# 与其他函数配合使用
data = [1, 2, 3]
multiplied = list(map(lambda x, y: x * y, data, itertools.repeat(3)))
print(multiplied)  # [3, 6, 9]

终止迭代器

chain() - 连接迭代器

将多个迭代器连接成一个。

import itertools

list1 = [1, 2, 3]
list2 = ['a', 'b']
list3 = [4, 5]

# 连接多个列表
combined = list(itertools.chain(list1, list2, list3))
print(combined)  # [1, 2, 3, 'a', 'b', 4, 5]

# 使用 chain.from_iterable() 处理嵌套列表
nested = [[1, 2], [3, 4], [5]]
flattened = list(itertools.chain.from_iterable(nested))
print(flattened)  # [1, 2, 3, 4, 5]

compress() - 条件过滤

根据选择器过滤数据。

import itertools

data = ['A', 'B', 'C', 'D', 'E']
selectors = [1, 0, 1, 0, 1]  # 1表示保留，0表示过滤

filtered = list(itertools.compress(data, selectors))
print(filtered)  # ['A', 'C', 'E']

takewhile() 和 dropwhile()

根据条件截取或跳过元素。

import itertools

numbers = [1, 3, 5, 8, 9, 11, 12]

# takewhile: 当条件为真时取元素，遇到假就停止
taken = list(itertools.takewhile(lambda x: x < 8, numbers))
print(taken)  # [1, 3, 5]

# dropwhile: 当条件为真时跳过元素，遇到假开始取后面所有元素
dropped = list(itertools.dropwhile(lambda x: x < 8, numbers))
print(dropped)  # [8, 9, 11, 12]

islice() - 切片迭代器

对迭代器进行切片操作。

import itertools

data = range(20)

# 取前5个元素
first_five = list(itertools.islice(data, 5))
print(first_five)  # [0, 1, 2, 3, 4]

# 从索引5开始取3个元素
middle = list(itertools.islice(data, 5, 8))
print(middle)  # [5, 6, 7]

# 带步长的切片
stepped = list(itertools.islice(data, 0, 10, 2))
print(stepped)  # [0, 2, 4, 6, 8]

组合迭代器

product() - 笛卡尔积

计算多个迭代器的笛卡尔积。

import itertools

colors = ['red', 'blue']
sizes = ['S', 'M', 'L']

# 计算所有组合
combinations = list(itertools.product(colors, sizes))
print(combinations)  
# [('red', 'S'), ('red', 'M'), ('red', 'L'), ('blue', 'S'), ('blue', 'M'), ('blue', 'L')]

permutations() - 排列

生成所有可能的排列。

import itertools

letters = ['A', 'B', 'C']

# 全排列
all_perms = list(itertools.permutations(letters))
print(all_perms)  
# [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

# 指定长度的排列
two_perms = list(itertools.permutations(letters, 2))
print(two_perms)  # [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

combinations() - 组合

生成指定长度的组合（不重复）。

import itertools

numbers = [1, 2, 3, 4]

# 2个元素的组合
pairs = list(itertools.combinations(numbers, 2))
print(pairs)  # [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]

# 可重复的组合
repeated_pairs = list(itertools.combinations_with_replacement([1, 2, 3], 2))
print(repeated_pairs)  # [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

分组函数

groupby() - 分组

将连续相同的元素分组。

import itertools

# 按值分组（需要先排序）
data = [1, 1, 2, 2, 2, 3, 1, 1]
for key, group in itertools.groupby(data):
    print(f"{key}: {list(group)}")
# 1: [1, 1]
# 2: [2, 2, 2] 
# 3: [3]
# 1: [1, 1]

# 按条件分组
words = ['apple', 'ant', 'bat', 'ball', 'cat']
words.sort()  # 先排序
for key, group in itertools.groupby(words, key=lambda x: x[0]):
    print(f"{key}: {list(group)}")
# a: ['ant', 'apple']
# b: ['ball', 'bat']
# c: ['cat']

itertools 模块提供了强大而高效的迭代工具，特别适合处理大数据集和复杂的迭代逻辑。合理使用这些函数可以让代码更简洁、更高效。