Double Iteration in List Comprehension

Viewed 294582

In Python you can have multiple iterators in a list comprehension, like

[(x,y) for x in a for y in b]

for some suitable sequences a and b. I'm aware of the nested loop semantics of Python's list comprehensions.

My question is: Can one iterator in the comprehension refer to the other? In other words: Could I have something like this:

[x for x in a for a in b]

where the current value of the outer loop is the iterator of the inner?

As an example, if I have a nested list:

a=[[1,2],[3,4]]

what would the list comprehension expression be to achieve this result:

[1,2,3,4]

?? (Please only list comprehension answers, since this is what I want to find out).

11 Answers

To answer your question with your own suggestion:

>>> [x for b in a for x in b] # Works fine

While you asked for list comprehension answers, let me also point out the excellent itertools.chain():

>>> from itertools import chain
>>> list(chain.from_iterable(a))
>>> list(chain(*a)) # If you're using python < 2.6

Gee, I guess I found the anwser: I was not taking care enough about which loop is inner and which is outer. The list comprehension should be like:

[x for b in a for x in b]

to get the desired result, and yes, one current value can be the iterator for the next loop.

Order of iterators may seem counter-intuitive.

Take for example: [str(x) for i in range(3) for x in foo(i)]

Let's decompose it:

def foo(i):
    return i, i + 0.5

[str(x)
    for i in range(3)
        for x in foo(i)
]

# is same as
for i in range(3):
    for x in foo(i):
        yield str(x)

This memory technic helps me a lot:

[ <RETURNED_VALUE> <OUTER_LOOP1> <INNER_LOOP2> <INNER_LOOP3> ... <OPTIONAL_IF> ]

And now you can think about Return + Outer-loop as the only Right Order

Knowing above, the order in list comprehensive even for 3 loops seem easy:


c=[111, 222, 333]
b=[11, 22, 33]
a=[1, 2, 3]

print(
  [
    (i, j, k)                            # <RETURNED_VALUE> 
    for i in a for j in b for k in c     # in order: loop1, loop2, loop3
    if i < 2 and j < 20 and k < 200      # <OPTIONAL_IF>
  ]
)
[(1, 11, 111)]

because the above is just a:

for i in a:                         # outer loop1 GOES SECOND
  for j in b:                       # inner loop2 GOES THIRD
    for k in c:                     # inner loop3 GOES FOURTH
      if i < 2 and j < 20 and k < 200:
        print((i, j, k))            # returned value GOES FIRST

for iterating one nested list/structure, technic is the same: for a from the question:

a = [[1,2],[3,4]]
[i2    for i1 in a      for i2 in i1]
which return [1, 2, 3, 4]

for one another nested level

a = [[[1, 2], [3, 4]], [[5, 6], [7, 8, 9]], [[10]]]
[i3    for i1 in a      for i2 in i1     for i3 in i2]
which return [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

and so on

I could never write double list comprehension on my first attempt. Reading into PEP202, it turns out the reason is that it was implemented in the opposite way you would read it in English. The good news is that it is a logically sound implementation, so once you understand the structure, it's very easy to get right.

Let a, b, c, d be successively nested objects. For me, the intuitive way to extend list comprehension would mimic English:

# works
[f(b) for b in a]
# does not work
[f(c) for c in b for b in a]
[f(c) for c in g(b) for b in a]
[f(d) for d in c for c in b for b in a]

In other words, you'd be reading from the bottom up, i.e.

# wrong logic
(((d for d in c) for c in b) for b in a)

However this is not how Python implements nested lists. Instead, the implementation treats the first chunk as completely separate, and then chains the fors and ins in a single block from the top down (instead of bottom up), i.e.

# right logic
d: (for b in a, for c in b, for d in c)

Note that the deepest nested level (for d in c) is farthest from the final object in the list (d). The reason for this comes from Guido himself:

The form [... for x... for y...] nests, with the last index varying fastest, just like nested for loops.

Using Skam's text example, this becomes even more clear:

# word: for sentence in text, for word in sentence
[word for sentence in text for word in sentence]

# letter: for sentence in text, for word in sentence, for letter in word
[letter for sentence in text for word in sentence for letter in word]

# letter:
#     for sentence in text if len(sentence) > 2, 
#     for word in sentence[0], 
#     for letter in word if letter.isvowel()
[letter for sentence in text if len(sentence) > 2 for word in sentence[0] for letter in word if letter.isvowel()]

This flatten_nlevel function calls recursively the nested list1 to covert to one level. Try this out

def flatten_nlevel(list1, flat_list):
    for sublist in list1:
        if isinstance(sublist, type(list)):        
            flatten_nlevel(sublist, flat_list)
        else:
            flat_list.append(sublist)

list1 = [1,[1,[2,3,[4,6]],4],5]

items = []
flatten_nlevel(list1,items)
print(items)

output:

[1, 1, 2, 3, 4, 6, 4, 5]
Related