摘要: 获取生成器长度，但不改变生成器状态

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

考虑一个简单的场景，有一个数组 nums，问其中小于 k 的元素个数有多少。这个问题本身很简单，直接线性枚举即可解决。

1
2
3
4

ans = 0
for x in nums:
    if x < k:
        ans += 1

如果想用函数式算法，可以用 Python 中的 filter 函数：

1

filtered = filter(lambda x: x < k, nums)

filtered 是一个 filter 类型的对象，它的行为与 generator 对象极为相似，并且两者之间可以互相转换。

首先他们都是迭代器，也就是实现了 __iter__() 和 __next__() 方法，这意味着可以通过 for 循环遍历它们，也可以用 next() 获取下一个元素。

其次两者都是惰性求值的，也就是不会在创建时立即计算所有值，而是在需要时才生成下一个值，这使得它们在处理大量数据时比较高效，不会一次性占用大量内存。这里注意 filter 对象虽然惰性求值，但其内存占用还要取决于被过滤的可迭代对象。

最后两者都是单次遍历的，即只能遍历一次，直至耗尽，此后就无法再次使用。如果需要再次遍历，需要重新创建一个新的生成器或 filter 对象。

generator 对象是 Python 中的基础特性，可以通过生成器表达式创建，或者通过生成器函数（使用 yield）创建。例如上面的 filter 如果改为生成器的话可以像下面这样写：

• 生成器函数

1
2
3
4
5
6

def my_generator():
    for x in nums:
        if x < k: 
            yield x

generator = my_generator()

• 生成器表达式

1

generator = (x for x in nums if x < k)

下面考虑上述函数式算法中的 filtered 对象（相当于生成器对象），如何获取其长度。当然最直接的是将其转换为列表再通过 len() 获取，这种方法会消耗额外内存。

1

n = len(list(filtered))

另一个很直接的办法是通过遍历生成器来计数，不会占用额外内存，

1
2
3

n = 0
for _ in filtered:
    n += 1

注意上述两种方法都会耗尽生成器，后续生成器就不能再用了。如果算法需要二次遍历，重新创建一个生成器即可。