在Python 3中使用itertools.product代替双嵌套for循环

Question

以下代码有效,但显示详细.

def gen(l):
    for x in range(l[0]):
        for y in range(l[1]):
            for z in range(l[2]):
                yield [x, y, z]
l = [1, 2, 3]
print(list(gen(l)))

>>>[[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [0, 1, 2]]

我的目的是用itertools.product减少LOC.这就是我想出来的.

from itertools import product
def gen(l):
    for x, y, z in product(map(range, l)):
        yield [x, y, z]
l = [1, 2, 3]
print(list(gen(l)))

ValueError: not enough values to unpack (expected 3, got 1)

是否有不同的方法来使用itertools.product以便有足够的值来解压缩？

Answer 1

您需要将map迭代器的元素product分别传递给*:

for x, y, z in product(*map(range, l))

顺便说一句,通过另一个map调用,您可以保存另一行,跳过Python生成器的开销,并在C中完成所有工作:

def gen(l):
    return map(list, product(*map(range, l)))