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

  • A+
Category:Languages

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.

 


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

Comment

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen: