假设我有一个任意的可迭代对象——例如,一个遍历文件行的生成器和
yield
s匹配正则表达式的行。
如果我不关心元素本身,我如何计算该可迭代对象中的项目数?
调用 Python 2 中的
itertools.imap()
或 Python 3 中的 map()
可以替换为等效的生成器表达式:
sum(1 for dummy in it)
这也使用了惰性生成器,因此它避免了在内存中具体化所有迭代器元素的完整列表。
当可迭代对象可能很长时,该方法比
sum(1 for i in it)
有意义地快(当可迭代对象很短时,它不会有意义地变慢),同时保持固定的内存开销行为(与len(list(it))
不同)以避免交换抖动和较大输入的重新分配开销:
# On Python 2 only, get zip that lazily generates results instead of returning list
from future_builtins import zip
from collections import deque
from itertools import count
# Avoid constructing a deque each time, reduces fixed overhead enough
# that this beats the sum solution for all but length 0-1 inputs
consumeall = deque(maxlen=0).extend
def ilen(it):
# Make a stateful counting iterator
cnt = count()
# zip it with the input iterator, then drain until input exhausted at C level
consumeall(zip(it, cnt)) # cnt must be second zip arg to avoid advancing too far
# Since count 0 based, the next value is the count
return next(cnt)
像
len(list(it))
一样,它在CPython上用C代码执行循环(deque
,count
和zip
都是用C实现的);避免每次循环执行字节码通常是 CPython 性能的关键。
想出公平的测试用例来比较性能是非常困难的(
list
使用 __length_hint__
的作弊,它不太可能用于任意输入迭代,itertools
不提供 __length_hint__
的函数通常有特殊的操作模式,当每个循环返回的值在请求下一个值之前被释放/释放时工作得更快,deque
和 maxlen=0
会做)。我使用的测试用例是使用 Python 3.3+ 的 itertools
:创建一个生成器函数,该函数接受一个输入并返回一个缺少特殊
__length_hint__
返回容器优化或
yield from
的 C 级生成器
def no_opt_iter(it):
yield from it
然后使用
ipython
%timeit
魔法(用不同的常数代替100):
>>> %%timeit fakeinput = (0,) * 100
... ilen(no_opt_iter(fakeinput))
当输入不够大以至于
len(list(it))
会导致内存问题时,在运行 Python 3.9 x64 的 Linux 机器上,无论输入长度如何,我的解决方案比 def ilen(it): return len(list(it))
花费大约 50% 的时间。
对于最小的输入,加载/调用
consumeall
/zip
/count
/next
的设置成本意味着这种方式比 def ilen(it): sum(1 for _ in it)
花费的时间无限长(在我的机器上大约多 40 ns输入长度为 0,比简单的sum
方法增加了 10%),但是当您输入长度为 2 时,成本是相等的,并且在长度为 30 左右的某个地方,与实际工作相比,初始开销不明显; sum
方法大约需要 50% 的时间。
基本上,如果内存使用很重要或输入没有限制大小并且您更关心速度而不是简洁,请使用此解决方案。如果输入有界且较小,
len(list(it))
可能是最好的,如果它们没有界,但简单/简洁很重要,您可以使用 sum(1 for _ in it)
.
一个简短的方法是:
def ilen(it):
return len(list(it))
请注意,如果您要生成lot元素(例如,数万个或更多),那么将它们放入列表中可能会成为性能问题。然而,这是对大多数情况下性能无关紧要的想法的简单表达。
more_itertools
是一个实现了ilen
工具的第三方库。 pip install more_itertools
import more_itertools as mit
mit.ilen(x for x in range(10))
# 10
len(list(it))
虽然,如果是无限生成器,它可以挂掉。
为此,我喜欢 cardinality 包,它非常轻量级,并尝试根据可迭代对象使用尽可能快的实现。
用法:
>>> import cardinality
>>> cardinality.count([1, 2, 3])
3
>>> cardinality.count(i for i in range(500))
500
>>> def gen():
... yield 'hello'
... yield 'world'
>>> cardinality.count(gen())
2
这些将是我的选择之一或另一个:
print(len([*gen]))
print(len(list(gen)))
如果您想在其他地方使用 iterable 并知道消耗了多少元素,您可以创建一个简单的包装类:
from collections.abc import Iterable, Iterator
from typing import Generic, TypeVar
_T = TypeVar("_T")
class IterCounter(Generic[_T]):
"""Iterator that keeps count of the consumed elements"""
def __init__(self, iterable: Iterable[_T]) -> None:
self._iterator = iter(iterable)
self.count = 0
def __iter__(self) -> Iterator[_T]:
return self
def __next__(self) -> _T:
element = next(self._iterator)
self.count += 1
return element
counter = IterCounter(range(5))
print(counter.count) # 0
print(list(counter)) # [0, 1, 2, 3, 4]
print(counter.count) # 5