确保参数可以迭代两次

小能豆

确保参数可以迭代两次

假设我有以下功能：

def print_twice(x):
    for i in x: print(i)
    for i in x: print(i)

当我跑步时：

print_twice([1,2,3])

或者：

print_twice((1,2,3))

我得到了预期的结果：数字 1,2,3 被打印两次。

但是当我跑步时：

print_twice(zip([1,2,3],[4,5,6]))

对 (1,4),(2,5),(3,6) 仅打印一次。可能是因为zip返回的生成器在一次传递后终止。

我如何修改该函数print_twice以便它能正确处理所有输入？

我可以在函数开头插入一行：x = list(x)。但如果 x 已经是列表、元组、范围或任何其他可以多次迭代的迭代器，那么这可能效率低下。有没有更有效的解决方案？

阅读 13

2024-12-18

共1个答案

小能豆

你描述的问题是由于 zip 返回的对象是一个 迭代器，而不是像列表或元组那样可以多次迭代。迭代器一旦被遍历，就会耗尽，无法再次遍历。

更高效的解决方案

我们可以检查输入对象是否是迭代器（即单次迭代后会耗尽的对象），然后仅在需要时将其转为列表。这样可以避免不必要的复制，提升效率。

实现

使用 collections.abc.Iterable 和 iter、next 方法来区分迭代器和可多次迭代的对象：

from collections.abc import Iterable

def print_twice(x):
    # 如果 x 是单次可消耗的迭代器，转为列表
    if not isinstance(x, Iterable) or iter(x) is x:
        x = list(x)

    for i in x: 
        print(i)
    for i in x: 
        print(i)

解释

isinstance(x, Iterable):
检查 x 是否是一个可迭代对象（例如，列表、元组、生成器等）。
如果 x 不是可迭代对象，我们不能对其进行迭代，函数将抛出错误。
iter(x) is x:
检查 x 是否是迭代器。
对于迭代器，iter(x) 返回自身（即 x），而对于列表、元组等多次可迭代对象，iter(x) 返回一个新的迭代器对象。
list(x):
仅当 x 是迭代器时，将其转为列表以支持多次遍历。

测试代码

print_twice([1, 2, 3])                # 多次遍历的列表
print("---")
print_twice((1, 2, 3))                # 多次遍历的元组
print("---")
print_twice(zip([1, 2, 3], [4, 5, 6])) # 单次遍历的 zip 对象
print("---")
print_twice(range(3))                 # 多次遍历的 range
print("---")
print_twice(iter([1, 2, 3]))          # 单次遍历的列表迭代器

输出

1
2
3
1
2
3
---
1
2
3
1
2
3
---
(1, 4)
(2, 5)
(3, 6)
(1, 4)
(2, 5)
(3, 6)
---
0
1
2
0
1
2
---
1
2
3
1
2
3

性能优化的说明

只有当输入是迭代器时，才会将其转为列表。对于像列表、元组、range 等多次可迭代对象，这样避免了额外的内存开销和计算。
这种方法在通用性和效率之间找到了平衡。

2024-12-18