Использование itertools.product вместо двойного вложенного цикла for в Python 3

Следующий код работает, но кажется многословным.

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

Мое намерение было сократить LOC с itertools.product. Вот что я придумал.

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, чтобы было достаточно значений для распаковки?

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)))
Другие вопросы по тегам