Python itertools.tee() 生成多相同迭代器

说明

《Python 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

itertools.tee() 是 Python itertools 模块中的一个函数，用于将一个可迭代对象分成多个独立的迭代器。这个函数非常有用，当你需要在多个地方同时迭代相同的数据时，避免多次遍历原始数据。itertools.tee() 可以将单个可迭代对象复制成多个独立的迭代器。每个迭代器都可以独立地从头开始迭代，而不会影响其他迭代器的状态。

语法

它的语法为 itertools.tee(iterable, n=2)，它从一个可迭代对象中返回 n 个独立的迭代器。

参数和返回值：

iterable：要拆分的可迭代对象，如列表、元组、字符串等。
n：要生成的迭代器数量，默认为 2。
返回值：返回一个元组，包含 n 个独立的迭代器，每个迭代器都可以独立地迭代原始数据。

下面的 Python 代码能帮助解释 tee 做了什么（尽管实际的实现更复杂，而且仅使用了一个底层的 FIFO 队列）。

大致相当于：

def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # when the local deque is empty
                try:
                    newval = next(it)   # fetch a new value and
                except StopIteration:
                    return
                for d in deques:        # load it to all the deques
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

一旦 tee() 实施了一次分裂，原有的 iterable 不应再被使用；否则 tee 对象无法得知 iterable 可能已向后迭代。

tee 迭代器不是线程安全的。当同时使用由同一个 tee() 调用所返回的迭代器时可能引发 RuntimeError，即使原本的 iterable 是线程安全的。

该迭代工具可能需要相当大的辅助存储空间（这取决于要保存多少临时数据）。通常，如果一个迭代器在另一个迭代器开始之前就要使用大部份或全部数据，使用 list() 会比 tee() 更快。

案例：

import itertools

t = itertools.tee(['a', 'b', 'c', 'd'], 3)

[*t]
'''
[<itertools._tee at 0x7f9d31802780>,
 <itertools._tee at 0x7f9d31800c00>,
 <itertools._tee at 0x7f9d31800400>]
'''

for i in t:
    print([*i])

'''
['a', 'b', 'c', 'd']
['a', 'b', 'c', 'd']
['a', 'b', 'c', 'd']
'''

示例

基本用法

import itertools

data = [1, 2, 3, 4, 5]

# 将 data 分成两个独立的迭代器
it1, it2 = itertools.tee(data)

# 迭代第一个迭代器
for item in it1:
    print(item, end=' ')  # 输出 1 2 3 4 5
print()

# 迭代第二个迭代器
for item in it2:
    print(item, end=' ')  # 输出 1 2 3 4 5
print()

生成多个迭代器

import itertools

data = [1, 2, 3, 4, 5]

# 将 data 分成三个独立的迭代器
it1, it2, it3 = itertools.tee(data, 3)

# 迭代第一个迭代器
print(list(it1))  # 输出 [1, 2, 3, 4, 5]

# 迭代第二个迭代器
print(list(it2))  # 输出 [1, 2, 3, 4, 5]

# 迭代第三个迭代器
print(list(it3))  # 输出 [1, 2, 3, 4, 5]

应用场景

多次迭代：当需要多次迭代相同的数据时，可以使用 itertools.tee() 生成多个独立的迭代器。
并行处理：在并行处理任务时，可以使用 itertools.tee() 将数据分发给多个处理器。

实用案例：同时处理数据的不同部分

假设我们有一个数据列表，我们希望同时处理其不同部分，例如计算其总和和乘积。

示例数据

data = [1, 2, 3, 4, 5]

目标：计算数据的总和和乘积。

实现代码：

import itertools
import functools
import operator

data = [1, 2, 3, 4, 5]

# 使用 itertools.tee 生成两个独立的迭代器
it1, it2 = itertools.tee(data)

# 计算总和
sum_result = sum(it1)
print(f"Sum: {sum_result}")  # 输出 Sum: 15

# 计算乘积
product_result = functools.reduce(operator.mul, it2, 1)
print(f"Product: {product_result}")  # 输出 Product: 120

解释：

导入 itertools 模块：使用 import itertools 导入 itertools 模块。
使用 itertools.tee() 生成两个独立的迭代器：itertools.tee(data) 将数据列表 data 拆分成两个独立的迭代器 it1 和 it2。
计算总和：使用 sum(it1) 计算第一个迭代器 it1 中所有元素的总和。
计算乘积：使用 functools.reduce(operator.mul, it2, 1) 计算第二个迭代器 it2 中所有元素的乘积。
打印结果：分别打印总和和乘积的结果。

这个实用案例展示了如何使用 itertools.tee() 生成多个独立的迭代器，从而可以同时对数据进行不同的处理操作，非常适用于需要多次迭代相同数据或并行处理的场景。

参考

https://docs.python.org/zh-cn/3/library/itertools.html#itertools.tee

< takewhile() 按条件终止迭代 Python 标准库 itertools 为高效循环而创建迭代器的函数 zip_longest() 最长拉链 >

更新时间：2024-07-04 08:15:49 标签：python itertools 迭代器