How do loops work in list assignments in python?

Question

Consider the following command:

elite_states  = [s for i in range(len(states_batch)) 
                if rewards_batch[i]>=reward_threshold for s in states_batch[i]]

I found that is equivalent to the following loop:

_l=[]
for i in range(len(states_batch)):
  for s in states_batch[i]:
    if rewards_batch[i]>=reward_threshold:
      _l.append(s)

However I don't get how the loop after s in the first command becomes the outer loop in its equivalent. I want to understand the format of the command so I get how it works!

It may seem backwards at first, but having the last loop in the list comprehension be the inner loop is so `i` is in scope. Try ordering them the other way around, the use of `i` would come before its declaration. — Mario Ishac, Jun 06 '20 at 18:25
https://stackoverflow.com/questions/3633140/nested-for-loops-using-list-comprehension — AMC, Jun 06 '20 at 18:28
if you think about it they _(the correct version of the loop as stated by @ForceBru)_ are in the same order: outerloop -> if -> innerloop, the only thing that changes is the position of the result 's' — Hoxha Alban, Jun 06 '20 at 18:34

ForceBru · Accepted Answer · 2020-06-06T18:38:21.987

I found that is equivalent to the following loop:

It isn't. List comprehensions are read in the same way as ordinary nested loops. First you loop over range(len(states_batch)), then, inside this loop, you check if rewards_batch[i]>=reward_threshold, next, inside this condition, the last loop is ran.

So the "expanded" version looks like this:

_l=[]
for i in range(len(states_batch)):
  if rewards_batch[i]>=reward_threshold:  # `if` and the inner `for` are switched
    for s in states_batch[i]:
      _l.append(s)

In code, this is extremely similar to the comprehension:

_l = [
  s
  for i in range(len(states_batch))
  if rewards_batch[i]>=reward_threshold
  for s in states_batch[i]
]

The only difference is that loops in comprehensions are written sequentially (one after the other), but in the expanded version they become nested one inside the other, the last loop in the comprehension being nested the deepest.

How do loops work in list assignments in python?

1 Answers1