List comprehensions

What you will be able to do

After this lesson you can

• turn a loop that builds a list into a list comprehension,
• filter with if, and choose a value with if ... else,
• write comprehensions with two loops, and
• decide when a plain loop is clearer.

A loop that builds a list

You have written this loop many times: make an empty list, loop, and append.

squares = []
for n in numbers:
    squares.append(n * n)

Python has a short way to write exactly this: a list comprehension.

Predict first. What will each line show? Then run it.

The result is a new list. The long loop and [n * n for n in numbers] give the same list, as the first line of output shows.

A useful way to read one aloud: "the square of n, for each n in numbers, if n is even".

Choosing a value: if ... else

There are two different ifs, and they go in two different places.

• A filter if goes at the end: [n for n in numbers if n % 2 == 0] keeps only the even numbers. Items that fail are left out.
• A choice if ... else goes at the front, as part of the expression: ["even" if n % 2 == 0 else "odd" for n in numbers] gives one value for every item. Nothing is left out.

A common mistake is to put else at the end. [n for n in numbers if n > 0 else 0] is a SyntaxError. The choice must come first: [n if n > 0 else 0 for n in numbers].

Two loops, and tables

A comprehension can also hold another comprehension in its expression. [[row * col for col in range(1, 4)] for row in range(1, 4)] makes a table: a list of rows, where each row is itself made by a comprehension. The multiplication table of 3 by 3 is [[1, 2, 3], [2, 4, 6], [3, 6, 9]].

To flatten a table into one list, put the loops in the order that you would write them as nested loops: [item for row in table for item in row].

The last example transposes a matrix: the rows become columns. [[row[i] for row in matrix] for i in range(3)] builds, for each position i, the list of the i-th items of all the rows. The matrix [[1, 2, 3], [4, 5, 6]] becomes [[1, 4], [2, 5], [3, 6]].

Read a comprehension from the middle: for each x in range(n), if the test passes, put x * x in the list.

When not to use one

Rules of thumb.

• If it does not fit comfortably in one line, use a loop.
• If you need to explain it, use a loop.
• Never use a comprehension for its side effects. [print(n) for n in numbers] builds a useless list of None. Use a plain for loop to print.
• A comprehension is faster than a loop with append, but the difference matters only for big lists. Choose for clarity first.

Common mistakes

• Putting else at the end. A SyntaxError. The choice goes at the front.
• Forgetting the brackets, and making a generator instead (the next lessons).
• A comprehension that is too long. Use a loop.
• Reversing the order of the two loops in a flatten. Write them as you would nest for statements.
• Using the variable of a comprehension outside it. It exists only inside.

Practice

Start with two quick questions. Each New numbers gives you a fresh one, so you can practise as long as you like.

Then the labs and the problems:

1. Warm-up. Open the lab Comprehension drills. You write four comprehensions, and the checker confirms each result.
2. Build. Open the problem Loop to comprehension and the problem Comprehension to loop.
3. Stretch. Open the problem A table and its transpose and write the transposed table of a 3 by 2 matrix.

Remember these

• [expression for item in iterable if condition] builds a new list.
• A filter if goes at the end. A choice if ... else goes at the front.
• Several for parts are nested loops, the first one outermost.
• Use a plain loop when the comprehension is hard to read, or for side effects.

Check yourself

Try each one on your own first. If you are stuck, open a hint. Then check with Show the answer.

1. What is [n * 2 for n in [1, 2, 3] if n > 1]?

2. Where does else go in a comprehension?

3. What does [x for row in [[1, 2], [3]] for x in row] give?