python - 如何检查生成器是否至少有 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

您可以 zip带有 range(n) 的迭代器，然后迭代并对元素进行计数:

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]