Enumerating a tuple of indices with itertools.product

Question

I am simultaneously iterating over multiple lists and want my generator to yield both the element and its index. If I had two lists, I would use a nested for loop:

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)

However, since I have more than two lists, the nested solution quickly becomes illegible. I would normally use itertools.product to neatly get the Cartesian product of my lists, but this strategy does not allow me to get the individual indices of elements in each list.

Here is what I have tried so far:

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

The output that I want is this:

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)

where the first, second, and third columns are the indices of elements from the respective list. Is there clean way of doing this that is still legible when there are a large number of lists?

Answer 1

You can use zip and product again, in a function:

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

For example:

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

For python2:

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