使用 itertools.product 枚举索引元组

Question

我同时迭代多个列表，并希望我的生成器生成元素及其索引。如果我有两个列表，我会使用嵌套的 for 循环：

for i_idx, i_val in enumerate(list_0):
    for j_idx, j_val in enumerate(list_1):
        print(i_idx, i_val, j_idx, j_val)

但是，由于我有两个以上的列表，嵌套的解决方案很快变得难以辨认。我通常会使用 itertools.product 整齐地获得我的列表的笛卡尔乘积，但这种策略不允许我获得每个列表中元素的单独索引。

这是我迄今为止尝试过的：

>>> from itertools import product

>>> list_0 = [1,2]
>>> list_1 = [3,4]
>>> list_2 = [5,6]

>>> for idx, pair in enumerate(product(list_0, list_1, list_2)):
...    print(idx, pair)
0 (1, 3, 5)
1 (1, 3, 6)
2 (1, 4, 5)
3 (1, 4, 6)
4 (2, 3, 5)
5 (2, 3, 6)
6 (2, 4, 5)
7 (2, 4, 6)

我想要的输出是这样的：

0 0 0 (1, 3, 5)
0 0 1 (1, 3, 6)
0 1 0 (1, 4, 5)
0 1 1 (1, 4, 6)
1 0 0 (2, 3, 5)
1 0 1 (2, 3, 6)
1 1 0 (2, 4, 5)
1 1 1 (2, 4, 6)

其中第一、第二和第三列是相应列表中元素的索引。当有大量列表时，是否有干净的方法仍然清晰易读？

Answer 1

您可以zip在一个函数中再次使用和产品：

def enumerated_product(*args):
    yield from zip(product(*(range(len(x)) for x in args)), product(*args))

例如：

>>> for idx, pair in enumerated_product(list_0, list_1, list_2):
...     print(idx, pair)
...
(0, 0, 0) (1, 3, 5)
(0, 0, 1) (1, 3, 6)
(0, 1, 0) (1, 4, 5)
(0, 1, 1) (1, 4, 6)
(1, 0, 0) (2, 3, 5)
(1, 0, 1) (2, 3, 6)
(1, 1, 0) (2, 4, 5)
(1, 1, 1) (2, 4, 6)

对于python2：

def enumerated_product(*args):
    for e in zip(product(*(range(len(x)) for x in args)), product(*args)):
        yield e