如何高效获取Python generator生成器的首元素与尾元素

生成器属于惰性迭代器，不支持随机索引访问，所有取值操作都会消费迭代器，和可随机访问的列表有本质区别，具体获取方法如下：

获取首元素 #

仅需消费生成器的第一个元素，时间复杂度O(1)，是所有操作中效率最高的：

def get_first(gen, default=None):
    try:
        return next(gen)
    except StopIteration:
        # 生成器为空时返回默认值
        return default

注意：调用该方法后原生成器会跳过第一个元素，如果后续需要使用完整生成器序列，可提前用itertools.tee复制迭代器副本，或者先将生成器转为列表存储。

获取尾元素 #

生成器不支持倒序访问，必须遍历完全部元素才能拿到尾元素，时间复杂度O(n)，n为生成器元素总数：

def get_last(gen, default=None):
    last_val = default
    for item in gen:
        last_val = item
    return last_val

Python 3.10+ 也可以用functools.reduce简化写法：

from functools import reduce
def get_last(gen, default=None):
    return reduce(lambda _, x: x, gen, default)

同时获取首尾元素的优化方案 #

避免两次遍历生成器，一次遍历即可同时拿到首尾值：

def get_first_last(gen, default=None):
    first = default
    last = default
    it = iter(gen)
    # 先取首元素
    try:
        first = next(it)
    except StopIteration:
        return first, last
    # 遍历剩余元素拿尾值
    last = first
    for item in it:
        last = item
    return first, last

注意事项 #

• 不要对无限生成器执行获取尾元素的操作，会触发无限循环
• 如果生成器元素量极大，不建议转列表后再取值，会占用过高内存，优先用上述方法直接操作生成器
• 所有操作都会消费原生成器，需要保留原生成器的场景请提前用itertools.tee创建副本

内容的提问来源于stack exchange，提问作者mark12345