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

Answer 2

将您的代码与 Netwave 的 答案混合在一起，可以通过使用

product

、

enumerate

和

zip

提供稍微不同的方法。

您正在枚举产品，但您想要的是枚举的产品。

所以，这个：

for prod in product(enumerate(list_0), enumerate(list_1), enumerate(list_2)) 将以以下形式给出结果：

((idx_0, elem_0), (idx_1, elem_1), (idx_2, elem_2)) 通过

zip

ping 结果，您将获得所需的输出：

for prod in product(enumerate(list_0), enumerate(list_1), enumerate(list_2)):
    idx, elems = zip(*prod)
    print(idx, elems)

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

---

对于“自动化”

枚举，您可以执行以下操作：

def enumerated_lists(*args):
    return [enumerate(arg) for arg in args]

for prod in product(*enumerated_lists(list_0, list_1, list_2)):
    ...