如何检查生成器是否至少有N个元素

Question

Python没有在生成器上提供len()。什么是检查生成器是否包含至少N个元素的最有效方法。

一些测试用例：

>>> iter_len_at_least(range(3), 3) 
True
>>> iter_len_at_least(range(3), 4)
False

我天真的实施：

import itertools

def iter_len_at_least(i, n):
    return len(list(None for x in itertools.islice(i, 0, n))) == n

还有更多的pythonic方式来检查吗？

Answer 1

您可以使用range(n) zip迭代器，然后迭代并计算元素：

def iter_len_at_least(i, n):
    return sum(1 for _ in zip(range(n), i)) == n

请记住，这会占用迭代器的下一个n元素：

it = iter(range(5))
iter_len_at_least(it, 3)
print(list(it))
# output: [3, 4]

没有办法避免这种情况，但我们可以使用itertools.tee将迭代器拆分为两个并返回一个新的迭代器以及结果：

def iter_len_at_least(i, n):
    i, it = itertools.tee(i)
    return it, sum(1 for _ in zip(range(n), i)) == n

it = iter(range(5))
it, result = iter_len_at_least(it, 3)
print(result)
print(list(it))
# output:
# True
# [0, 1, 2, 3, 4]