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

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)))

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

Using itertools.product in place of double-nested for loop in Python 3

python

generator

cartesian-product

itertools

python-3.x