展开嵌套生成器表达式

11
我试图展平一个嵌套的生成器,但是我得到了一个意外的结果:
>>> g = ((3*i + j for j in range(3)) for i in range(3))
>>> list(itertools.chain(*g))
[6, 7, 8, 6, 7, 8, 6, 7, 8]

我希望结果看起来像这样:

[0, 1, 2, 3, 4, 5, 6, 7, 8]

我认为我得到意外的结果是因为内部生成器直到外部生成器已经被迭代过并设置i为2之后才被评估。我可以通过强制使用列表推导式而不是生成器表达式来评估内部生成器来拼凑出一个解决方案:

>>> g = ([3*i + j for j in range(3)] for i in range(3))
>>> list(itertools.chain(*g))
[0, 1, 2, 3, 4, 5, 6, 7, 8]
理想情况下,我希望有一个完全懒惰的解决方案,不会强制评估嵌套元素的内部直到它们被使用。
是否有一种方法可以展开任意深度的嵌套生成器表达式(也许使用除了 itertools.chain 之外的其他东西)?
编辑:
不,我的问题不是Variable Scope In Generators In Classes的重复。我老实说无法理解这两个问题如何相关。也许管理员可以解释一下为什么他认为这是一个重复的问题。
此外,对于我的问题,两个答案都是正确的,因为它们可以用来编写正确展开嵌套生成器的函数。
def flattened1(iterable):
    iter1, iter2 = itertools.tee(iterable)
    if isinstance(next(iter1), collections.Iterable):
        return flattened1(x for y in iter2 for x in y)
    else:
        return iter2

def flattened2(iterable):
    iter1, iter2 = itertools.tee(iterable)
    if isinstance(next(iter1), collections.Iterable):
        return flattened2(itertools.chain.from_iterable(iter2))
    else:
        return iter2

就我使用 timeit 的观察,它们的性能表现是相同的。

>>> timeit(test1, setup1, number=1000000)
18.173431718023494
>>> timeit(test2, setup2, number=1000000)
17.854709611972794

就样式而言,我也不确定哪个更好,因为x for y in iter2 for x in y有点晦涩难懂,但可以说比itertools.chain.from_iterable(iter2)更优雅。欢迎提供意见。

遗憾的是,我只能标记其中一个同样好的答案为正确。

3个回答

12

您可以使用chain.from_iterable代替使用chain(*g)

>>> g = ((3*i + j for j in range(3)) for i in range(3))
>>> list(itertools.chain(*g))
[6, 7, 8, 6, 7, 8, 6, 7, 8]
>>> g = ((3*i + j for j in range(3)) for i in range(3))
>>> list(itertools.chain.from_iterable(g))
[0, 1, 2, 3, 4, 5, 6, 7, 8]

8
这个怎么样:
[x for y in g for x in y]

产生的结果是:
[0, 1, 2, 3, 4, 5, 6, 7, 8]

3
我猜你已经有了答案,但这里提供另一种看法。问题在于每个内部生成器被创建时,产生值的表达式都会封闭外部变量i,因此即使第一个内部生成器开始生成值时,它也使用i的“当前”值。如果外部生成器已经完全消耗,那么i的值将为i=2(在chain(*g)调用中的参数求值后,在实际调用chain之前正是这种情况)。以下阴险的技巧将解决这个问题:
g = ((3*i1 + j for i1 in [i] for j in range(3)) for i in range(3))

请注意,这些内部生成器并没有捕获变量 i ,因为for子句在生成器创建时就已经被求值,所以单例列表[i]被求值,并且其值在面对i变量的进一步更改时“冻结”。
相比from_iterable答案,这种方法有一个优点,即它更通用,如果您想在chain.from_iterable调用之外使用它-- 它将始终产生“正确”的内部生成器,无论在内部生成器使用之前外部生成器是否被部分或全部消耗。例如,在以下代码中:
g = ((3*i1 + j for i1 in [i] for j in range(3)) for i in range(3))
g1 = next(g)
g2 = next(g)
g3 = next(g)

你可以插入以下行:
list(g1)
list(g2)
list(g3)

在定义内部生成器后的任何时间和顺序,您将获得正确的结果。

+1 这真的很巧妙,很酷,但它并不能帮助我扁平化我所使用的嵌套生成器。 - castle-bravo
1
不,我想你是对的!但是,如果您以这样的方式构建嵌套生成器,使得内部生成器的值取决于外部生成器的评估时间,那么它就是一个编程定时炸弹——一些小的代码重构将导致一切崩溃。无论您最终如何展平g,重新设计内部生成器以避免不良变量捕获可能是个好主意。 - K. A. Buhr
1
你说的一切都是正确的,但我的(未明确说明的)目标不是在所有地方使用像range()这样的生成器(range()可以很好地生成扁平化的数字范围),而是让其他程序员能够使用生成器定义多维数组。如果他们认为他们的多维数组看起来像[[0,1],[2,3]],而实际上它看起来像[[2,3],[2,3]]——而且没有引发任何错误——那么我就搞砸了。 - castle-bravo

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接