Python 模块 itertools 包含 20 个工具,每个 Python 开发者都应该了解。
本文将 itertools 模块中的迭代器分为 5 类,以便于学习,同时我们还列出了3个最常用的迭代器。
分类 | 迭代器 |
---|---|
重塑迭代器 | batched , chain* , groupby , islice , pairwise* |
过滤迭代器 | compress , dropwhile , filterfalse , takewhile |
组合迭代器 | combinations , combinations_with_replacement , permutations , product* |
无限迭代器 | count , cycle , repeat |
补充其他工具的迭代器 | accumulate , starmap , zip_longest |
除了上表中列出的 19 个迭代器外,itertools 模块还提供了一个非常强大的函数 tee,tee()
函数虽然没有被明确归类,但它非常有用。
tee()
函数返回n个独立的迭代器,这些迭代器都与原始迭代器相关联,但可以独立地消耗。这对于并行处理同一数据流的不同部分很有用,同时避免了多次遍历整个数据集的开销。
itertools
模块中的 3 个最有用的迭代器
根据个人经验和使用场景,itertools
模块中最常用的三个工具是 product
、chain
和 pairwise
。
product
展平嵌套循环
product
迭代器是一种组合迭代器,当你需要展平一系列嵌套的 for 循环时非常有用。举个典型的例子,遍历二维网格的嵌套循环可以用 product
重写为一个单一的循环。
例如,当我们有两个或多个独立的嵌套 for 循环时,如下所示:
for x in range(width):
for y in range(height):
# 执行操作...
我们可以使用 product
将它们重塑为一个单独的循环:
from itertools import product
for x, y in product(range(width), range(height)):
# 执行操作...
这不仅让代码看起来更整洁,而且也使得扩展到更高维度更加容易。例如,如果需要处理三维或更高维度的空间,只需向product
函数添加更多的范围即可,而不需要增加额外的嵌套层次。
下面是使用itertools.product
遍历一个三维网格的例子:
from itertools import product
width = 10
height = 10
depth = 10
for x, y, z in product(range(width), range(height), range(depth)):
# 在这里执行操作,现在我们有了 x, y 和 z 的所有可能组合
print(f"Processing point ({x}, {y}, {z})")
在这个例子中,product
接受三个范围作为输入,并生成一个包含所有(x, y, z)三元组的迭代器,其中x从0到9,y从0到9,z也是从0到9。
这使得代码可以轻易地处理更高维度的问题,同时保持简洁和可读性。
chain
整合多个序列
在处理多个序列时,itertools.chain
提供了一种优雅的方法,使你能够像操作单一序列那样轻松遍历多个序列。这种方法尤其适用于需要合并多个列表、元组或其他序列类型的场景,避免了显式连接操作,从而保持代码的简洁与高效。
传统的做法是通过加号(+)将序列串联起来形成一个新的序列,如下所示:
# 常见的序列合并模式:
first_list = [1, 2, 3]
second_list = [4, 5, 6]
full_list = first_list + second_list # 结果为 [1, 2, 3, 4, 5, 6]
for element in full_list:
# 对每个元素执行操作
然而,通过itertools.chain
,我们可以跳过序列的直接拼接步骤,直接在一个循环中访问所有元素:
from itertools import chain
first_list = [1, 2, 3]
second_list = [4, 5, 6]
for element in chain(first_list, second_list):
# 对每个元素执行操作,无需预先创建完整序列
这种技巧在处理生成器时尤为有用,因为生成器不能直接相加:
first_gen = (x ** 2 for x in range(3)) # 生成器表达式
second_gen = (x ** 3 for x in range(3))
for value in chain(first_gen, second_gen):
print(value, end=" ") # 输出:0 1 4 0 1 8
虽然可以考虑将生成器转换为列表再进行拼接,但这样做既浪费资源又不适用于处理无限迭代器的情况。
chain
还提供了一个便捷的辅助构造函数chain.from_iterable
,它能将一个包含多个可迭代对象的容器展平,生成一个单一的迭代器。这一特性在需要遍历嵌套列表时特别有用:
nested_lists = [[1, 2, 3], [4], [], [5, 6]]
flattened_list = list(chain.from_iterable(nested_lists))
print(flattened_list) # 输出:[1, 2, 3, 4, 5, 6]
# 直接遍历元素,无需显式转换为列表:
for value in chain.from_iterable(nested_lists):
print(value, end=" ") # 输出:1 2 3 4 5 6
通过itertools.chain
及其from_iterable
方法,你可以有效地简化代码,提升性能,同时保持代码的可读性和可维护性。
pairwise
高效构建重叠元素对
itertools.pairwise
是一个功能强大的迭代器,它接受任意可迭代对象,并巧妙地生成连续元素的重叠对,这一操作对于处理序列中的相邻元素对比或连接十分有效。相较于手动实现,pairwise
提供了更高效且通用的解决方案,尤其是在处理大型数据集时优势显著。
通常,开发者可能会采用zip(my_list[:-1], my_list[1:])
的模式来生成元素对。然而,这种方法存在两个潜在缺点:
- 性能成本:对于大规模的可迭代对象,使用切片操作来生成子列表可能相当耗时且占用内存。
- 适用性限制:并非所有的可迭代对象都支持切片操作,如生成器或文件流等。
使用itertools.pairwise
则能规避上述问题,提供更为简洁和优化的实现方式。例如,在以下场景中,我们希望让列表中的每个人向其后一位朋友打招呼:
names = ["Harry", "Anne", "George"]
# 传统方式
for left, right in zip(names[:-1], names[1:]):
print(f"{left} says hi to {right}")
# 使用pairwise
from itertools import pairwise
for left, right in pairwise(names):
print(f"{left} says hi to {right}")
无论采用哪种方式,输出结果都将一致:
Harry says hi to Anne
Anne says hi to George
然而,通过itertools.pairwise
,我们不仅省去了zip
和切片
的繁琐步骤,还获得了更高的性能和更广泛的适用性,使代码更加优雅且高效。
重塑迭代器
在Python中,有一些迭代器可以用于重新排列和重新格式化输入数据,从而生成不同格式的输出。这些迭代器提供了简洁且高效的方法来处理数据。
签名 | 文档 | 简要说明 |
---|---|---|
batched(iterable, n) | Docs | 从给定的可迭代对象生成长度为 n 的元组,直到耗尽。最后一个元组可能少于 n 个元素。 |
chain(*iterables) | Docs | 将多个可迭代对象连接成一个单一的可迭代对象。 |
chain.from_iterable(iterable) | Docs | 将一个可迭代对象的可迭代元素展平成一个单一的可迭代对象。 |
islice(iterable, stop) | Docs | 从给定的可迭代对象中切片前 stop 个元素。类似于 lst[:stop]。 |
islice(iterable, start, stop[, step]) | Docs | 从给定的可迭代对象中切片,跳过前 start 个元素,并返回到 stop 位置之间的元素,仅返回每 step 个元素。类似于 lst[start:stop:step]。 |
groupby(iterable, key=None) | Docs | 创建子迭代器,用于分组连续值,使得对于每个组,函数 key 返回相同的值。 |
pairwise(iterable) | Docs | 生成重叠的连续元素对。类似于 zip(lst[:-1], lst[1:])。 |
下面是每种重塑迭代器的简单示例。
1. batched(iterable, n)
将可迭代对象分割成长度为n
的元组,直到耗尽。最后一个元组可能少于n
个元素。
import itertools
def batched(iterable, n):
it = iter(iterable)
while True:
batch = tuple(itertools.islice(it, n))
if not batch:
break
yield batch
# 示例
for batch in batched(range(10), 3):
print(batch)
输出:
(0, 1, 2)
(3, 4, 5)
(6, 7, 8)
(9,)
2. chain(*iterables)
将多个可迭代对象连接成一个单一的可迭代对象。
from itertools import chain
# 示例
lst1 = [1, 2, 3]
lst2 = ['a', 'b', 'c']
for item in chain(lst1, lst2):
print(item)
输出:
1
2
3
a
b
c
3. chain.from_iterable(iterable)
将一个可迭代对象的可迭代元素展平成一个单一的可迭代对象。
from itertools import chain
# 示例
nested_list = [[1, 2, 3], ['a', 'b', 'c'], [4, 5]]
for item in chain.from_iterable(nested_list):
print(item)
输出:
1
2
3
a
b
c
4
5
4. islice(iterable, stop)
从给定的可迭代对象中切片前stop
个元素。类似于lst[:stop]
。
from itertools import islice
# 示例
for item in islice(range(10), 5):
print(item)
输出:
0
1
2
3
4
5. islice(iterable, start, stop[, step])
从给定的可迭代对象中切片,跳过前start
个元素,并返回到stop
位置之间的元素,仅返回每step
个元素。类似于lst[start:stop:step]
。
from itertools import islice
# 示例
for item in islice(range(10), 2, 8, 2):
print(item)
输出:
2
4
6
6. groupby(iterable, key=None)
创建子迭代器,用于分组连续值,使得对于每个组,函数key
返回相同的值。
from itertools import groupby
# 示例
data = [('a', 1), ('a', 2), ('b', 3), ('b', 4), ('a', 5)]
for key, group in groupby(data, lambda x: x[0]):
print(key, list(group))
输出:
a [('a', 1), ('a', 2)]
b [('b', 3), ('b', 4)]
a [('a', 5)]
7. pairwise(iterable)
生成重叠的连续元素对。类似于zip(lst[:-1], lst[1:])
。
from itertools import tee
def pairwise(iterable):
a, b = tee(iterable)
next(b, None)
return zip(a, b)
# 示例
for a, b in pairwise([1, 2, 3, 4, 5]):
print(a, b)
输出:
1 2
2 3
3 4
4 5
这些迭代器为数据处理提供了强大的工具,可以根据需要重构和重新格式化数据。每个迭代器都有其特定的用途,适合不同的应用场景。
过滤迭代器
过滤迭代器接受一个可迭代对象和一个谓词函数(predicate),并生成原始可迭代对象的一个子集。下面是每个迭代器的简单示例。
签名 | 文档 | 简要说明 |
---|---|---|
compress(data, selectors) | Docs | 生成data 中对应selectors 为真值的元素。 |
dropwhile(predicate, iterable) | Docs | 丢弃iterable 中满足给定谓词的首个连续元素序列。 |
filterfalse(predicate, iterable) | Docs | 内置filter 的补充。生成iterable 中不满足给定函数的元素。 |
takewhile(predicate, iterable) | Docs | dropwhile 的补充。生成iterable 中满足给定函数的首个连续值序列。 |
下面是每种过滤迭代器的简单示例。
示例
compress(data, selectors)
这个函数从data
中选择元素,选择依据是selectors
中的布尔值。如果selectors
中的值为True
,则对应的data
中的元素会被选中。
from itertools import compress
data = 'ABCDEF'
selectors = [1, 0, 1, 0, 1, 0]
result = compress(data, selectors)
print(list(result))
输出:
['A', 'C', 'E']
dropwhile(predicate, iterable)
这个函数会丢弃iterable
中最初连续的一系列元素,直到predicate
首次返回False
。之后的元素都会被保留。
from itertools import dropwhile
data = [1, 4, 6, 4, 1]
predicate = lambda x: x < 5
result = dropwhile(predicate, data)
print(list(result))
输出:
[6, 4, 1]
filterfalse(predicate, iterable)
这是filter()
函数的补充,它生成iterable
中不满足predicate
的所有元素。
from itertools import filterfalse
data = [1, 4, 6, 4, 1]
predicate = lambda x: x < 5
result = filterfalse(predicate, data)
print(list(result))
输出:
[6]
takewhile(predicate, iterable)
这个函数是dropwhile
的补充,它生成iterable
中满足predicate
的首个连续值序列。
from itertools import takewhile
data = [1, 4, 6, 4, 1]
predicate = lambda x: x < 5
result = takewhile(predicate, data)
print(list(result))
输出:
[1, 4]
这些过滤迭代器为数据处理提供了强大的工具,可以根据需要选择性地生成和过滤数据。每个迭代器都有其特定的用途,适合不同的应用场景。
组合迭代器
组合迭代器是Python的itertools
模块中的一类函数,用于在不同的可迭代对象之间创建各种组合。以下是几种常见的组合迭代器及其用途:
签名 | 文档 | 简要说明 |
---|---|---|
combinations(iterable, r) | Docs | 从给定的可迭代对象中生成长度为 r 的元组,其中元素按其原始位置排序。 |
combinations_with_replacement(iterable, r) | Docs | 与 combinations 相同,但每个值可以任意多次重复。 |
permutations(iterable, r=None) | Docs | 生成给定可迭代对象中 r 个元素的所有排列。 |
product(*iterables, repeat=1) | Docs | 生成元组,组合来自所有给定可迭代对象的所有元素。可迭代对象可以重复任意次数。 |
需要注意的是在处理组合和排列时,itertools
模块中的迭代器如permutations
, combinations
, 和 combinations_with_replacement
都是基于位置而不是基于元素的实际值来生成结果的。这意味着即使集合中包含重复的元素,只要它们在生成的序列中的位置不同,也会被视为不同的结果。
示例
combinations(iterable, r)
返回从输入可迭代对象中取r长度的所有可能组合,组合不考虑顺序。
from itertools import combinations
data = ['A', 'B', 'C']
result = combinations(data, 2)
print(list(result))
输出:
[('A', 'B'), ('A', 'C'), ('B', 'C')]
combinations_with_replacement(iterable, r)
与combinations
类似,但是元素可以重复选取。
from itertools import combinations_with_replacement
data = ['A', 'B', 'C']
result = combinations_with_replacement(data, 2)
print(list(result))
输出:
[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]
permutations(iterable, r=None)
返回从输入可迭代对象中取r长度的所有可能排列。
from itertools import permutations
data = ['A', 'B', 'C']
result = permutations(data, 2)
print(list(result))
输出:
[('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
product(*iterables, repeat=1)
计算笛卡尔积,即从一个或多个可迭代对象中取所有可能的元素组合。如果提供了repeat
参数,表示重复使用相同的可迭代对象。
from itertools import product
data1 = [1, 2]
data2 = ['A', 'B']
result = product(data1, data2)
print(list(result))
输出:
[(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]
这些迭代器可以非常高效地生成大量组合或排列,尤其适用于算法和数据分析等领域。使用itertools
模块的组合迭代器可以替代传统的嵌套循环,提高代码的可读性和效率。
无限迭代器
itertools
模块中的无限迭代器可以用于生成无限序列,这对于某些算法和数据流处理非常有用。
签名 | 文档 | 简要说明 |
---|---|---|
count(start=0, step=1) | Docs | 这个函数会从start 值开始,每次迭代增加step 值,从而产生一个无限递增的整数序列。 |
cycle(iterable) | Docs | 这个函数会无限次地重复iterable 中的元素。 |
repeat(object[, times]) | Docs | 这个函数会无限次地重复object ,或者如果指定了times 参数,则重复times 次。 |
示例
count(start=0, step=1)
生成从 start
开始的无限等差数列。
from itertools import count
for i in count(10, 2):
if i > 20:
break
print(i, end=" ") # 输出: 10 12 14 16 18 20
cycle(iterable)
这个函数会无限次地重复iterable
中的元素。
from itertools import cycle
count = 0
for item in cycle('AB'):
if count >= 6:
break
print(item, end=" ") # 输出: A B A B A B
count += 1
repeat(object[, times])
这个函数会无限次地重复object
,或者如果指定了times
参数,则重复times
次。
from itertools import repeat
for item in repeat('Hello', 3):
print(item, end=" ") # 输出: Hello Hello Hello
这些无限迭代器在需要生成连续数据流或重复执行特定操作的场景下非常有用。通过与其他迭代器结合使用,可以实现复杂的数据处理和算法设计。
补充其他工具的迭代器
itertools
模块还提供了几个其他的迭代器,它们补充了Python标准库中已有的功能。下面列出的是几种常用的补充迭代器:
签名 | 补充 | 文档 | 简要说明 |
---|---|---|---|
accumulate(iterable[, function, *, initial=None]) | functools.reduce | Docs | 这个迭代器会累积iterable 中的元素,通常用于计算累加或累乘等操作。 |
starmap(function, iterable) | map | Docs | 类似于map ,但是starmap 将iterable 中的每个元素(通常是元组或列表)解包并传递给function 。 |
zip_longest(*iterables, fillvalue=None) | zip | Docs | 类似于内置的zip 函数,但是当最短的可迭代对象耗尽时,zip_longest 会继续直到最长的可迭代对象结束,并用fillvalue 填充缺失的值。 |
示例
accumulate(iterable[, function, *, initial=None])
这个迭代器会累积iterable
中的元素,通常用于计算累加或累乘等操作。
from itertools import accumulate
import operator
data = [1, 2, 3, 4]
result = accumulate(data, operator.add)
print(list(result)) # 输出: [1, 3, 6, 10]
# 使用乘法函数
result = accumulate(data, operator.mul)
print(list(result)) # 输出: [1, 2, 6, 24]
starmap(function, iterable)
类似于map
,但是starmap
将iterable
中的每个元素(通常是元组或列表)解包并传递给function
。
from itertools import starmap
data = [(2, 3), (4, 5), (6, 7)]
result = starmap(lambda x, y: x * y, data)
print(list(result)) # 输出: [6, 20, 42]
zip_longest(*iterables, fillvalue=None)
类似于内置的zip
函数,但是当最短的可迭代对象耗尽时,zip_longest
会继续直到最长的可迭代对象结束,并用fillvalue
填充缺失的值。
from itertools import zip_longest
data1 = [1, 2, 3]
data2 = ['a', 'b']
result = zip_longest(data1, data2, fillvalue='?')
print(list(result)) # 输出: [(1, 'a'), (2, 'b'), (3, '?')]
这些迭代器与其他工具配合使用,为数据处理提供了更多的灵活性和功能。
accumulate
和starmap
使得处理复杂的数据流变得简单,而zip_longest
则帮助处理长度不一的序列组合问题。
每个迭代器都有其特定的用途,适合不同的应用场景。
tee 函数
tee
函数是Python的itertools
模块中一个非常有用且有趣的功能。迭代器的基本概念是一次性的——一旦被遍历,就不能再次从头开始遍历。但是,tee
函数打破了这个常规,允许你从同一个迭代器创建多个独立的副本,每个副本都可以独立地遍历原始迭代器的数据,而不影响其他副本的状态。
下面是如何使用tee
函数的一个简单示例:
from itertools import tee
# 创建一个简单的迭代器
original_iterable = iter([1, 2, 3, 4])
# 使用tee函数创建两个独立的迭代器副本
iter1, iter2 = tee(original_iterable, 2)
# 遍历第一个迭代器
for item in iter1:
print(item)
# 再次遍历第二个迭代器
for item in iter2:
print(item)
在这个例子中,尽管iter1
已经被完全遍历了,iter2
仍然能够从头开始遍历整个序列。这是因为tee
函数实际上是在内部缓存了迭代器的数据,以便每个生成的迭代器都能访问完整序列。
值得注意的是,由于tee
需要缓存数据,所以当处理大数据集时要谨慎使用,因为这可能会导致较高的内存消耗。另外,如果原始迭代器是无限的,那么tee
函数将无法完成缓存过程,因为没有终点可以确定缓存何时结束。
在实际应用中,tee
函数可以用于并行处理相同数据的不同部分,或者在需要多次遍历相同数据流但不想重复生成数据的情况下。