赛博红兔的科技博客

大家好，我是赛博红兔。你有没有遇到过这样的情况：想要遍历多个列表的组合，却发现自己写的嵌套循环又长又复杂？或者想要高效地生成排列和组合，但写出的代码效率不太理想？今天，我们将深入探讨Python中的itertools模块。这个模块包含了一系列常用的迭代器，以及用于组合多个迭代器的函数，能帮我们轻松地生成排列、组合、无限序列还有批量分组等等。

1. count方法 —— 无限计数器

简单来说，迭代器可以理解为一组可以被遍历的数据。itertools里的count方法能够提供一个无限递增的迭代器。默认情况下，它从 0 开始，每次递增 1，并且会一直运行下去，直到你手动终止它。注意啊，这段代码会不断输出 0、1、2、3、4……，永远不会停止！如果你不小心运行了它，可能会让你的电脑崩溃。因此，在运行时请务必小心，通常我们不会这样直接遍历它。我们可以使用 Python 内置的 next() 函数来获取 count() 生成的下一个值，而不会陷入无限循环。我们也可以通过设置参数让 count() 从任意数值开始，并按指定步长递增。

import itertools

counter = itertools.count()
for num in counter:
    print(num)

print(next(counter))
print(next(counter))
print(next(counter))

counter = itertools.count(start=5, step=5)
print(next(counter))
print(next(counter))
print(next(counter))

2. cycle方法—— 无限循环迭代器

cycle() 允许我们无限循环一个序列，比如这里列表里的红、绿、蓝就在不断地循环生成。

colors = itertools.cycle(["红", "绿", "蓝"])
print(next(colors))
print(next(colors))
print(next(colors))
print(next(colors))

3. repeat方法—— 重复生成相同的值

repeat方法会无限重复某个值，如果我们只想重复有限次，可以设置 times 参数，比如这里不断重复数字2三次。repeat() 通常用于 map() 和 zip() 这样的函数，例如计算 1-10 的平方。

repeat_2 = itertools.repeat(2, times=3)
print(list(repeat_2))

squares = map(pow, range(10), itertools.repeat(2))
print(list(squares))

4. combinations方法 —— 组合

在处理数据时，我们常常需要计算不同元素的组合。组合是指从一个序列中选择一定数量的元素，但顺序无关，即 (A, B) 和 (B, A) 视为相同的组合。例如，假设我们有一个包含 A, B, C, D 的列表，我们需要找出每两个一对元素的组合。可以看到，每一对元素的顺序是不重要的，例如 (A, B) 和 (B, A) 只出现了一次。组合通常用在像选拔队员、彩票号码生成等场景。比如说，在扑克游戏中，不管你是先抽到 A 还是 K，最终的牌型是一样的。如果你希望组合里的元素可以重复可以使用 combinations_with_replacement()。这在类似于投骰子或密码生成的场景下非常有用。

letters = ["A", "B", "C", "D"] 
result = itertools.combinations(letters, 2) 
print(list(result))

result = itertools.combinations_with_replacement(letters, 2)
print(list(result))

5. permutations方法——排列

排列与组合类似，但顺序是重要的。也就是说，(A, B) 和 (B, A) 视为不同的排列。在这段代码里，我们可以看到所有的排列里，每个元素对的顺序是可以不同的 ，像(A, B) 和 (B, A) 都出现了。排列通常用于像竞赛排名、密码破解、搜索引擎优化等等场景。例如，在一场比赛中，选手 A 和 B 的名次不同，即 (A, B) 和 (B, A) 是两种不同的结果。如果你想要所有可能的排列，包括重复的元素，可以使用 product()。这在类似于密码破解或者坐标生成的场景非常有用，比如创建所有可能的 (x, y) 坐标场景。

letters = ["A", "B", "C", "D"]
result = itertools.permutations(letters, 2)
print(list(result))

result = itertools.product(letters, repeat=2)
print(list(result))

6. itertools.zip_longest方法 —— 处理不同长度的迭代器

在 Python 的内置 zip() 函数中，当其中最短的迭代器结束后，整个 zip() 也会终止，比如像这样。但如果我们想要继续匹配较长的迭代器，并用 None 填充缺失值，可以使用 itertools.zip_longest()。如果你想用 自定义填充值，可以使用 fillvalue 参数。

a = [1, 2, 3]
b = ['a', 'b']
print(list(zip(a, b)))

result = itertools.zip_longest(a, b)
print(list(result))

result = itertools.zip_longest(a, b, fillvalue="空")
print(list(result))

7. chain方法 —— 链接多个迭代器

如果你有多个可迭代对象（列表、元组、集合等），希望逐个遍历它们的时候，你可以使用 chain()。相比于直接用 + 拼接列表，chain() 更高效，因为它不会创建新的列表，而是直接逐个迭代元素。

list1 = ["A", "B", "C"]
list2 = [1, 2, 3]

result = itertools.chain(list1, list2)
print(list(result))

8. itertools.islice方法—— 高效切片

如果你想对迭代器进行切片（就是类似于列表里的切片），但又不希望把它转换成列表，你可以使用 islice()。这种方式比 list(numbers)[2:8:2] 这样的列表切片更高效，特别适用于大数据集。

numbers = range(10)

result = itertools.islice(numbers, 2, 8, 2) 
print(list(result))

9. groupby方法 —— 按键分组

如果你有一组字典数据，并希望按某个属性或者叫键来进行分组，你可以使用 groupby()。但要注意，groupby() 需要数据先排序，否则分组结果可能不正确。这个方法类似于SQL中的 GROUP BY，可以用于统计分析、数据整理等场景。

people = [
    {"name": "小明", "city": "北京"},
    {"name": "小红", "city": "北京"},
    {"name": "张伟", "city": "上海"},
    {"name": "李强", "city": "上海"},
    {"name": "王芳", "city": "北京"},
]

# 按 "city" 进行分组
people.sort(key=lambda x: x["city"])  # 必须先排序
grouped = itertools.groupby(people, key=lambda x: x["city"])

for key, group in grouped:
    print(key, list(group))

最后来总结一下，itertools 模块提供了一系列高效的迭代器工具，包括无限迭代、排列组合、迭代器操作等等，帮助我们更高效地处理数据，减少内存占用，并简化复杂的迭代逻辑。好啦，本频道将同步在B站和油管上更新，如果你喜欢我的内容，请不要忘记点赞、订阅和分享，这样就不会错过我更新的内容，欢迎在评论区提出你的想法和建议。今天就到这里，再见吧！