python - 使用 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