前言
大家好，我是海鸽。
我们知道，在Python中，迭代器是一个非常强大的利器，可以更高效地处理数据。然而，当处理大量数据时，迭代器的效率可能受到影响。
如果你还不清楚迭代器，你可以先温习下我的这篇文章。
为了解决这个问题，我们可以利用Python标准库中的itertools模块。itertools专门为处理迭代器而设计，提供了许多高效的函数，用于创建、组合和操作迭代器。每个函数都经过精心设计，以在内存使用和性能之间取得平衡，使其能够处理大型数据集。
一个常见的面试问题：
您有三张 20 美元的钞票、五张 10 美元的钞票、两张 5 美元的钞票和五张 1 美元的钞票。100 美元的钞票有多少种找零的方法？
要“暴力破解”这个问题，你只需开始列出从钱包中选择一张钞票的方法，检查其中是否有 100 美元的零钱，然后列出从钱包中选择两张钞票的方法，再次检查，等等等等。
但你是一名程序员，所以你自然希望自动化这个过程。恰好你又知道itertools这个库，因此你很快就写下下面的代码，然后就愉快地摸鱼去了~
import itertools
bills = [20, 20, 20, 10, 10, 10, 10, 10, 5, 5, 1, 1, 1, 1, 1]
makes_100 = []
for n in range(1, len(bills) + 1):
for combination in itertools.combinations(bills, n):
if sum(combination) == 100:
makes_100.append(combination)
接下来，我们将详细介绍一些常用的itertools函数，并通过示例演示它们的使用方法。
无穷迭代器
无限计数器：count(start=0, step=1)
count(start, step) 用于创建一个无限的迭代器，从 start 开始，每次增加 step。
相当于：
def count(start=0, step=1):
# count(10) --> 10 11 12 13 14 ...
# count(2.5, 0.5) --> 2.5 3.0 3.5 ...
n = start
while True:
yield n
n += step
因此：
import itertools
for num in itertools.count(1, 2):
if num > 10:
break
print(num) # 1 3 5 7 9
使用count()函数自定义my_enumerate函数，实现和enumerate一样的功能：
from itertools import count
def my_enumerate(iterable, start=0):
return zip(count(start), iter(iterable))
my_enumerate2 = lambda x, start=0: zip(count(start), x)
print(list(my_enumerate('word'))) # [(0, 'w'), (1, 'o'), (2, 'r'), (3, 'd')]
print(list(enumerate(iter('word')))) # [(0, 'w'), (1, 'o'), (2, 'r'), (3, 'd')]
for index, value in my_enumerate(["red", "green", "blue"]):
print(index, value)
# 0 red
# 1 green
# 2 blue
for index, value in my_enumerate2(["red", "green", "blue"]):
print(index, value)
# 0 red
# 1 green
# 2 blue
循环迭代：cycle(iterable)
cycle(iterable) 会无限重复迭代一个可迭代对象。
from itertools import cycle
colors = cycle(['red', 'green', 'blue'])
for _ in range(5):
print(next(colors)) # red green blue red green
重复生成指定元素：repeat(object[, times])
重复生成指定元素。
from itertools import repeat
for _ in repeat('Hello', 3):
print(_) # Hello Hello Hello
排列组合迭代器
笛卡尔积：product(*iterables, repeat=1)
计算多个可迭代对象的笛卡尔积。
from itertools import product
colors = ['red', 'green']
sizes = ['small', 'large']
result = product(colors, sizes)
print(list(result)) # [('red', 'small'), ('red', 'large'), ('green', 'small'), ('green', 'large')]
排列组合：permutations(iterable, r=None)
生成指定长度的所有可能排列，无序，无重复元素。
from itertools import permutations
items = ['A', 'B', 'C']
result = permutations(items, 2)
print(list(result)) # [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
组合生成：combinations(iterable, r)
生成指定长度的所有可能组合，有序，无重复元素。
from itertools import combinations
items = ['A', 'B', 'C']
result = combinations(items, 2)
print(list(result)) # [('A', 'B'), ('A', 'C'), ('B', 'C')]
组合生成（允许重复）：combinations_with_replacement(iterable, r)
返回长度为 r 的子序列，其中元素可能重复。
from itertools import combinations_with_replacement
items = ['A', 'B', 'C']
result = combinations_with_replacement(items, 2)
print(list(result)) # [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]
条件迭代器
链接迭代器：chain(*iterables)
将多个可迭代对象按顺序连接成一个序列。
from itertools import chain
list1 = [1, 2, 3]
tuple1 = ('a', 'b', 'c')
combined = chain(list1, tuple1)
print(list(combined)) # [1, 2, 3, 'a', 'b', 'c']
平铺迭代器：chain.from_iterable(iterable)
用于将嵌套的可迭代对象（例如嵌套的列表）平铺为单层的迭代器。这在处理嵌套结构时非常方便，避免了多层嵌套的循环。
from itertools import chain
flattened_strings = chain.from_iterable(['Hello', 'ab'])
print(list(flattened_strings)) # ['H', 'e', 'l', 'l', 'o', 'a', 'b']
nested_list = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flattened_iterable = chain.from_iterable(nested_list)
print(list(flattened_iterable)) # [1, 2, 3, 4, 5, 6, 7, 8, 9]
数据过滤：compress(data, selectors)
根据selectors筛选data中的元素。
from itertools import compress
data = ['a', 'b', 'c', 'd']
selectors = [True, False, 1, 0]
result = compress(data, selectors)
print(list(result)) # ['a', 'c']
返回不满足条件的元素：filterfalse(predicate, iterable)
创建一个过滤来自 iterable 中元素从而只返回其中 predicate 为假值的元素的迭代器。
from itertools import filterfalse
numbers = range(10)
result = filterfalse(lambda x: x % 2 == 0, numbers)
print(list(result)) # [1, 3, 5, 7, 9]
分组迭代器：groupby(iterable, key=None)
创建一个迭代器，它将可迭代对象的元素按照键分组。
import itertools
# 按照键进行分组
for key, group in itertools.groupby('AAAAABBCCCCCDDDCCCBBA'):
print(list(group))
# ['A', 'A', 'A', 'A', 'A']
# ['B', 'B']
# ['C', 'C', 'C', 'C', 'C']
# ['D', 'D', 'D']
# ['C', 'C', 'C']
# ['B', 'B']
# ['A']
# 按照国家分组
data = [('New York', 'US'), ("Shanghai", "China"),
("LA", 'US'), ("Shenzhen", "China")]
new_data = sorted(data, key=lambda x: x[1])
for country, group in itertools.groupby(new_data, lambda x: x[1]):
print(f"country: {country}, {list(group)}")
# country: China, [('Shanghai', 'China'), ('Shenzhen', 'China')]
# country: US, [('New York', 'US'), ('LA', 'US')]
迭代器切片：islice(iterable, stop)
对可迭代对象进行切片，类似于列表的切片操作。
from itertools import filterfalse
numbers = range(10)
result = filterfalse(lambda x: x % 2 == 0, numbers)
print(list(result)) # [1, 3, 5, 7, 9]
累积：accumulate(iterable, func=None)
对可迭代对象累积运算。
from itertools import accumulate
numbers = [1, 2, 3, 4, 5]
result = accumulate(numbers, lambda x, y: x + y)
print(list(result)) # [1, 3, 6, 10, 15]
result2 = accumulate(numbers, lambda x, y: x * y)
print(list(result2)) # [1, 2, 6, 24, 120]
result3 = accumulate(numbers)
print(list(result3)) # [1, 3, 6, 10, 15]
元素打包：zip_longest(*iterables, fillvalue=None)
将多个可迭代对象逐个元素打包，以最长的序列为准，缺失值用指定的填充值填充。
from itertools import zip_longest
list1 = [1, 2, 3]
tuple1 = ('a', 'b')
result = zip_longest(list1, tuple1, fillvalue='N/A')
print(list(result)) # [(1, 'a'), (2, 'b'), (3, 'N/A')]
result2 = zip_longest(list1, tuple1)
print(list(result2)) # [(1, 'a'), (2, 'b'), (3, None)]
复制迭代器：tee(iterable, n=2)
将一个可迭代对象复制为n个独立的迭代器。
from itertools import tee
# 不复制迭代器
numbers = [1, 2, 3]
numbers = (i * 2 for i in numbers)
print(list(numbers)) # [2, 4, 6]
# 重点： 生成器的最大的缺点之一就是：它会枯竭。当你完整遍历过它们后，之后的重复遍历就不能拿到任何新内容了。
print(list(numbers)) # []
# 使用tee复制迭代器
numbers = [1, 2, 3, 4, 5]
iter1, iter2 = tee(numbers, 2)
print(list(iter1)) # [1, 2, 3, 4, 5]
print(list(iter2)) # [1, 2, 3, 4, 5]
函数应用于可迭代对象的元素：starmap(function, iterable)
根据可迭代对象中的参数计算结果。
from itertools import starmap
def add(x, y):
return x + y
# 内置函数map示例, map函数可接收多个可迭代对象
numbers1 = [1, 2, 3]
numbers2 = [4, 5, 6]
result_map = map(add, numbers1, numbers2)
print(list(result_map)) # [5, 7, 9]
# itertools.starmap示例
number_pairs = [(1, 4), (2, 5), (3, 6)]
result_starmap = starmap(add, number_pairs)
print(list(result_starmap)) # [5, 7, 9]
条件终止获取：takewhile(predicate, iterable)
从可迭代对象中取元素，直到不满足条件为止，不再进一步检查。
from itertools import takewhile
numbers = [1, 3, 5, 2, 4, 6, 7]
result = takewhile(lambda x: x % 2 != 0, numbers)
print(list(result)) # [1, 3, 5]
条件开始获取：dropwhile(predicate, iterable)
从可迭代对象中取元素，直到不满足条件为止，不再继续检查。与takewhile(predicate, iterable)相反。
from itertools import dropwhile
numbers = [1, 3, 5, 2, 4, 6, 7]
result = dropwhile(lambda x: x % 2 != 0, numbers)
print(list(result)) # [2, 4, 6, 7]
迭代相邻的元素对： pairwise(iterable)
在可迭代对象中获取相邻的元素对。这在一些算法中很有用，比如计算相邻元素的差异或比较。
from itertools import pairwise
print(list(pairwise("ABCDEF"))) # [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F')]
分成批次（3.12 新版功能）：itertools.batched(iterable, n)
iterable 的长度为 n 元组形式的批次数据，最后一个批次可能短于 n。
from itertools import batched
flattened_data = ['roses', 'red', 'violets', 'blue', 'sugar', 'sweet']
unflattened = list(batched(flattened_data, 2))
print(unflattened) # [('roses', 'red'), ('violets', 'blue'), ('sugar', 'sweet')]
for batch in batched('ABCDEFG', 3):
print(batch)
# ('A', 'B', 'C')
# ('D', 'E', 'F')
# ('G',)
总结
itertools库对大多数人来说，是一个未被充分利用甚至很少听到的库，但它确实包含一些非常好用的函数，特别适用于数据科学、算法实现和其他需要高效处理迭代对象的领域。这些函数通常可以帮助我们将代码行数减少到仅一行，使我们的代码看起来更加简洁优雅！
参考
https://docs.python.org/3/library/itertools.html
https://realpython.com/python-itertools/#et-tu-brute-force
👨👩 朋友，希望本文对你有帮助～🌐
欢迎点赞 👍、在看 💙、关注 💡 三连支持一下～🎈
扫一扫
175
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
