如何用嵌套生成器与yield from保留嵌套结构实现元素过滤？

你现在需要处理嵌套列表a = [[1, 2], [3, 4, 5, 6], [7, 7, 8]]，要过滤每个子容器里的偶数，同时保留原有的嵌套结构，目标输出是[[2], [4, 6], [8]]。之前用列表推导的nested_filter能实现，但想换成两个生成器函数避免创建中间列表，不过用yield from的写法会把结构扁平化，那该怎么调整呢？

先回顾你现有的代码：

现有列表版实现（可保留结构） #

def nested_filter(obj):
    res = []
    for sub_obj in obj:
        even = [i for i in sub_obj if i % 2 == 0]
        res.append(even)
    return res

a = [[1, 2], [3, 4, 5, 6], [7, 7, 8]]
print(nested_filter(a))  # 输出: [[2], [4, 6], [8]]

问题：扁平化的生成器实现（非期望） #

不管是直接嵌套循环的生成器，还是用yield from调用子生成器的写法，都会把结果扁平化：

# 直接嵌套循环的扁平化生成器
def nested_yield(obj):
    """Will flatten the nested structure of obj."""
    for sub_obj in obj:
        for i in sub_obj:
            if i % 2 == 0:
                yield i

list(nested_yield(a))  # 输出: [2, 4, 6, 8]

# 用yield from的错误实现
def gen1(obj):
    for sub_obj in obj:
        yield from gen2(sub_obj)

def gen2(sub_obj):
    for i in sub_obj:
        if i % 2 == 0:
            yield i

list(gen1(a))  # 输出: [2, 4, 6, 8]

正确的嵌套生成器实现 #

要保留嵌套结构，关键是让外层生成器每次生成一个子生成器的迭代结果容器，而不是直接把内层元素平铺出来。这里提供两种实现方式：

方案1：拆分两个生成器函数 #

让内层生成器负责过滤子列表的偶数，外层生成器负责保留子结构，每次返回一个内层生成器实例：

def gen_inner(sub_obj):
    for i in sub_obj:
        if i % 2 == 0:
            yield i

def gen_outer(obj):
    for sub_obj in obj:
        # 每次yield一个内层生成器，保留子列表层级
        yield gen_inner(sub_obj)

a = [[1, 2], [3, 4, 5, 6], [7, 7, 8]]
# 转换为嵌套列表需要两层list转换
result = [list(sub_gen) for sub_gen in gen_outer(a)]
print(result)  # 输出: [[2], [4, 6], [8]]

方案2：用生成器表达式简化内层逻辑 #

如果不想单独写gen_inner，可以把内层过滤逻辑写成生成器表达式，直接在外层生成器中返回：

def nested_generator(obj):
    for sub_obj in obj:
        yield (i for i in sub_obj if i % 2 == 0)

# 同样通过两层转换得到嵌套列表
result = [list(sub_gen) for sub_gen in nested_generator(a)]
print(result)  # 输出: [[2], [4, 6], [8]]

为什么之前的yield from不行？ #

yield from gen2(sub_obj)会把gen2产生的每个元素直接 yield 到外层，相当于把内层生成器的元素“平铺”到外层输出里，自然就丢失了嵌套结构。而我们需要的是外层生成器每次输出一个“子迭代器”，而非子迭代器里的单个元素，这样就能保留原有的层级。

额外验证 #

如果直接迭代gen_outer(a)，你会得到三个生成器对象：

for sub_gen in gen_outer(a):
    print(sub_gen)
# 输出示例:
# <generator object gen_inner at 0x1023a5eb0>
# <generator object gen_inner at 0x1023a5f40>
# <generator object gen_inner at 0x1023a5fd0>

把每个子生成器转换成列表，就得到了我们想要的嵌套结构。

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