Build list using nested for loops with list comprehension

Question

I need to build a list using a list comprehension. This is basically what it has to do:

pattern = []
for c in range(3):
    for r in range(3):
        if r == c:
            pattern.append(a)
        else:
            pattern.append(b)

But this all somehow needs to be condensed to only one line! I have never used a list comprehension before, so please explain your solution.

Thank you!

Edit: If I wanted the new list to consist of sublists, could that be put into the list comprehension too? Above I used a range of 3, so in the produced list every sublist would consist of 3 elements, ie

pattern = [['a','b','b'],['b','a','b'],['b','b','a']]

Answer 1

The general form of the list comprehension can be understood like this

[(the task to be done) (how long it has to be done)]

We normally use for loop in the how long it has to be done part and the task to be done part can have if conditions as well. The important thing to be noted is, the task to be done part should return a valid value (even None is a valid value). So, you cannot use any of the Python statements ( return , in Python 2.x print , etc) in the list comprehension.

Now to answer your first question,

['a' if r == c else 'b' for c in range(3) for r in range(3)]
# ['a', 'b', 'b', 'b', 'a', 'b', 'b', 'b', 'a']

This exactly creates a list as you have shown in the for loop version.

'a' if r == c else 'b'

This is same as

if r == c:
    'a'
else:
    'b'

First, for c in range(3) will be executed and the list [0, 1, 2] will be generated, and then on every iteration for r in range(3) will be executed and the list [0, 1, 2] will be generated. On each iteration of r , the if condition we saw above will be executed and the result of that if..else will be used as the element of the new list being generated.

To answer your second question, you can very well use list comprehension.

Our basic understanding from the above example, is that list comprehension is to generate a list. Now, lets try and nest the list comprehensions (we are going to use list comprehension in the the task to be done part), like this

[['a' if r == c else 'b' for r in range(3)] for c in range(3)]
# [['a', 'b', 'b'], ['b', 'a', 'b'], ['b', 'b', 'a']]

Now, if you look at our nested list comprehension, first, for c in range(3) will be executed and then ['a' if r == c else 'b' for r in range(3)] part will be executed, which will generate the individual rows of the nested lists. The important thing to note here is, c is available within the nested list comprehension.