List comprehensions

Basics

List comprehensions are a Python feature that let us quickly write simple loops that and return the results to a list. This is a common pattern in programming, so it can save us a lot of keystrokes, but they are also often simpler to read for basic problems.

For example, let's say we have a list of numbers and we need to do some math (let's square each value) on each of those numbers and store the results in a list. Using a for loop the answer would look something like this.

In [1]:
numbers = [2.2, 7.3, 1, 0, 9.5, 6.2, 42]
results = []
for value in numbers:
    results.append(value ** 2)
results
Out[1]:
[4.840000000000001, 53.29, 1, 0, 90.25, 38.440000000000005, 1764]

Using a list comprehension we can right the middle 3 lines as a single line:

In [2]:
numbers = [2.2, 7.3, 1, 0, 9.5, 6.2, 42]
results = [value ** 2 for value in numbers]
results
Out[2]:
[4.840000000000001, 53.29, 1, 0, 90.25, 38.440000000000005, 1764]

Nested loops

List comprehensions can also handle nested for loops. So if we wanted to get the products of all of the possible combinations of numbers from 1 to 4 we could write a list comprehension like this:

In [3]:
multiplication_table = [a * b for a in range(1, 5) for b in range(1, 5)]
multiplication_table
Out[3]:
[1, 2, 3, 4, 2, 4, 6, 8, 3, 6, 9, 12, 4, 8, 12, 16]

The loops can interact just like in a normal nested for loop, so if we only wanted to calculate each unique combination once (i.e., only calculate 1 3, not also 3 1) we can modify the second for loop.

In [4]:
multiplication_table = [a * b for a in range(1, 5) for b in range(a, 5)]
multiplication_table
Out[4]:
[1, 2, 3, 4, 4, 6, 8, 9, 12, 16]

Conditionals

They can also handle conditions. So, if in the first example we only wanted, to square the number if it was positive we could check that inside the list comprehension.

In [5]:
numbers = [2.2, 7.3, -5, -2.2, 1, 0, 9.5, -7, 6.2, 42]
results = [value ** 2 for value in numbers if value > 0]
results
Out[5]:
[4.840000000000001, 53.29, 1, 90.25, 38.440000000000005, 1764]

Limits to readability

Simple list comprehensions are easier to read than their equivalent for loops, but this is only true up to a point. A general rule of thumb is to use a regular for loop if you have more than total for loops and conditions, or the if the conditions are complex.

Want to learn more?

Check out the Python documentation on list comprehensions.