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!

Answer 1

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.