How do I enumerate a list comprehension with two 'for's?

Question

I am trying to do

ls = [myfunc(a,b,i) for a in a_list for b in b_list]

but also pass in i into myfunc, which is an index starting at 0 and incrementing for each new element.

For example:

a_list = 'abc'
b_list = 'def'

should result in

ls = [myfunc('a','d',0),
      myfunc('a','e',1),
      myfunc('a','f',2),
      myfunc('b','d',3),
      myfunc('b','e',4),
      ...
      myfunc('c','f',8]

I know that I can use enumerate() for just the normal case, ie.

ls = [myfunc(a,i) for a,i in enumerate(a_list)]

But I can't figure out how to do it cleanly when there are two fors. I couldn't find this question posted previously either.

Answer 1

You are creating a Cartesian product over two lists, so use itertools.product() instead of a double for loop. This gives you a single iterable you can easily add enumerate() to:

from itertools import product

ls = [myfunc(a, b, i) for i, (a, b) in enumerate(product(a_list, b_list))]

For cases where you can't use product(), you'd put the multiple loops in a generator expression, then add enumerate() to that. Say you needed to filter some values of a_list:

gen = (a, b for a in a_list if some_filter(a) for b in b_list)
ls = [myfunc(a, b, i) for i, (a, b) in enumerate(gen)]

Another option is to add a separate counter; itertools.count() gives you a counter object that produces a new value with next():

from itertools import count

counter = count()
ls = [myfunc(a, b, next(counter)) 
      for a in a_list if some_filter(a)
      for b in b_list]

After all, in essence enumerate(iterable, start=0) is the equivalent of zip(itertools.count(start), iterable).