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.

Martijn Pieters · Accepted Answer · 2018-07-12T11:31:50.670

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

@tobias_k: that wasn't me. That was that guy over there with the paper bag over their head. — Martijn Pieters, Jul 12 '18 at 11:22
@PetterFriberg: *That's the one!* He keeps breaking into my house! They think I can't see them, but I can spot them when I walk past a mirror! — Martijn Pieters, Jul 12 '18 at 11:24

khelwood · Answer 2 · 2018-07-12T11:27:29.740

6

You can use enumerate on the sequence of pairs.

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

edited Jul 12 '18 at 11:27

answered Jul 12 '18 at 11:22

khelwood

55,782
14
81
108

I'll give you an upvote if you split that unreadable line into two readable ones (: Edit: That's not really what I meant, but it's readable enough I guess. – Aran-Fey Jul 12 '18 at 11:26
@Aran-Fey Sure. – khelwood Jul 12 '18 at 11:27

score 2 · Answer 3 · answered Jul 12 '18 at 11:24

2

For a simple nested loop, use itertools.product() as @Martijn suggested.

If the expression is more complex you can use itertools.count:

i_gen = itertools.count()
ls = [myfunc(a, b, next(i_gen)) for a in a_list for b in b_list]

answered Jul 12 '18 at 11:24

Duncan

92,073
11
122
156

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

3 Answers3