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?

Primogenial

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

Supplicatory

Answer 2

A mix of your code and the answer by Netwave can give a slightly different approach by using product, enumerate and zip.

You were enumerating the product, but what you want is the product of the enumerations.

So, this:

for prod in product(enumerate(list_0), enumerate(list_1), enumerate(list_2))

Will give results in the form of:

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

And by zipping the result you get your desired output:

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)

---

And to "automate" the enumeration you can do something like:

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

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

Laina