List Comprehensions

Basic Comprehensions

Let’s say we have a list of numbers and we want to double each number. With what we have learned so far, our code would look something like this:

>>> my_favorite_numbers = [1, 1, 2, 3, 5, 8, 13]
>>> doubled_numbers = []
>>> for n in my_favorite_numbers:
...     doubled_numbers.append(n * 2)
...
>>> doubled_numbers
[2, 2, 4, 6, 10, 16, 26]

In Python there is a shorter syntax for this. We can write the code to create our doubled_numbers list in only one line:

>>> doubled_numbers = [n * 2 for n in my_favorite_numbers]
>>> doubled_numbers
[2, 2, 4, 6, 10, 16, 26]

This is called a list comprehension. List comprehensions provide convenient shorthand for creating lists from other lists or iterables. We can put any expression that makes a new object inside of the first part of a comprehension. A simple list comprehension is handy when we want to initialize a list:

>>> [0 for i in range(10)]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Let’s create a list of number tuples where the second item of the tuple is the square of the first:

>>> squares = [(x, x ** 2) for x in range(1, 11)]
>>> squares
[(1, 1), (2, 4), (3, 9), (4, 16), (5, 25), (6, 36), (7, 49), (8, 64), (9, 81), (10, 100)]

Now let’s combine iterations of two lists of strings to create a list of strings of every combination of the lists:

>>> colors = ["red", "green", "blue", "yellow"]
>>> things = ["ball", "car", "house"]
>>> combos = [" ".join([color, item]) for color in colors for item in things]
>>> combos
['red ball', 'red car', 'red house', 'green ball', 'green car', 'green house', 'blue ball', 'blue car', 'blue house', 'yellow ball', 'yellow car', 'yellow house']

We can call functions or methods within the list comprehension:

>>> caps = [color.upper() for color in colors]
>>> caps
['RED', 'GREEN', 'BLUE', 'YELLOW']

We could also use a list comprehension to get specific digits in a string:

>>> number = 4321
>>> digits = [int(d) for d in str(number)]
>>> print(digits)
[4, 3, 2, 1]

List comprehensions can perform their operations on any iterable. For example we could loop over a tuple of numbers and double each:

>>> numbers = (12, 4, 56, 5, -1, 38)
>>> doubles = [x * 2 for x in numbers]
>>> doubles
[24, 8, 112, 10, -2, 76]
>>> tuple(doubles)
(24, 8, 112, 10, -2, 76)

We could also loop over all keys in a dictionary:

>>> colors = {"purple": 0, "aqua": 1, "blue": 2}
>>> stretched_colors = ["_".join(word) for word in colors]
>>> stretched_colors
['p_u_r_p_l_e', 'a_q_u_a', 'b_l_u_e']

Of course, when we use an unordered collection type for the iterable, there is no guarantee about the order of the items returned, but you will always get all the items.

We can nest list comprehensions to make more complicated lists. Let’s create a matrix of incremental numbers, then create the transpose of the matrix:

>>> matrix = [[row * 3 + incr for incr in range(1, 4)] for row in range(4)]
>>> matrix
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
>>> transpose = [[row[i] for row in matrix] for i in range(3)]
>>> transpose
[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]]

In the real world, built-in functions are preferred, so, as we saw previously, we can transpose the matrix using zip:

>>> transpose = list(zip(*matrix))
>>> transpose
[(1, 4, 7, 10), (2, 5, 8, 11), (3, 6, 9, 12)]

But zip gives you immutable tuples for rows, instead of lists. What if you need to be able to modify the matrix?

You can use a list comprehension on the transposed matrix:

>>> transpose = [list(row) for row in transpose]
>>> transpose
[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]]

Or you can do it all in one command at the same time as the transpose:

>>> transpose = [list(row) for row in list(zip(*matrix))]
>>> transpose
[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]]

Conditional Filters

A powerful feature of list comprehensions is the ability to use a conditional if clause to add a filter to the iterable. Say we want a list of cubes of all perfect squares up through 100:

>>> [x ** 3 for x in range(101) if (x ** 0.5).is_integer()]
[0, 1, 64, 729, 4096, 15625, 46656, 117649, 262144, 531441, 1000000]

Let’s make a list comprehension that gets all numbers from a list that are greater than zero:

>>> nums = [4, -1, 7, 9, 34, 0, -4, 3]
>>> new_nums = [x for x in nums if x > 0]
>>> new_nums
[4, 7, 9, 34, 3]

We could also check which pets are too numerous, given a dictionary of pets and quantities:

>>> num_pets = {'cats' : 6, 'dogs' : 4, 'hamsters' : 7, 'birds' : 3}
>>> too_many = [pets for pets, num in num_pets.items() if num > 4]
>>> too_many
['cats', 'hamsters']

When do we use them?

List comprehensions are a tool. Like all tools, you need to be able to identify opportunities to use them.

You can use list comprehensions whenever you see a “for loop” that loops over an iterable, transforming each item and adding it to a list.

Take this function:

>>> def square_all(numbers):
...     squared_numbers = []
...     for n in numbers:
...         squared_numbers.append(n * n)
...     return squared_numbers
...
>>> square_all([1, 2, 3, 4])
[1, 4, 9, 16]
>>> square_all((0, 1, 2))
[0, 1, 4]

You can see there is a “for loop” making one list (or any other iterable) into a new list.

We can rewrite this as a list comprehension:

>>> def square_all(numbers):
...     return [n * n for n in numbers]
...
>>> square_all([1, 2, 3, 4])
[1, 4, 9, 16]
>>> square_all((0, 1, 2))
[0, 1, 4]

We can also use list comprehensions whenever we see a “for loop” that also excludes some values, using an “if statement” to filter out values that don’t meet a condition.

Take this function:

>>> def only_truthy(things):
...     truthy_things = []
...     for x in things:
...         if x:
...             truthy_things.append(x)
...     return truthy_things
...
>>> only_truthy([1, 0, 2, False, "hello", True, ""])
[1, 2, 'hello', True]

That “if statement” can be transformed into the condition statement in a list comprehension.

We can rewrite this as a list comprehension like this:

>>> def only_truthy(things):
...     return [x for x in things if x]
...
>>> only_truthy([1, 0, 2, False, "hello", True, ""])
[1, 2, 'hello', True]

That looks a little strange with all those “x” variables in there. We have “x for x” because we’re not actually modifying each item in this list. We’re just filtering them out.

We can also modify the values at the same time:

>>> nums = [4, -1, 7, 9, 34, 0, -4, 3]
>>> new_nums = [x * 3 for x in nums if x > 0]
>>> new_nums
[12, 21, 27, 102, 9]

Readability

To make your comprehensions more readable, I recommend always breaking them over multiple lines of code.

Sometimes it’s easier to start a complicated comprehension by writing a for loop first and then copy-pasting your way from a loop to a comprehension.

Start with a loop that we want to make into a comprehension:

pet_counts = {'cats': 6, 'dogs': 4, 'hamsters': 7, 'birds': 3}
too_many = []
for pet, num in pet_counts.items():
    if num > 4:
        too_many.append(pet)

Let’s copy-paste our way into a comprehension:

too_many = [
    pet  # The thing we were appending to the list
    for pet, num in pet_counts.items()  # Looping through the items
    if num > 4  # Expression to decide whether to include this element.
]

Comprehensions can always be written over multiple lines and doing so often improves readability.

Comprehension Exercises

Tip

Most of these you have seen before, and solved with for loops. You should use at least one list comprehension in each of these exercises! Each one can be tested by running python test.py function_name from your exercise directory.

Hint

If you get stuck for a minute or more, try searching Google or using help.

If you’re stuck for more than a few minutes, some of these links might be helpful for some of the exercises below:

• Copy-pasting from a loop to a comprehension

• Turning nested loops into comprehensions

• Breaking strings up by a separator

• Looping over multiple iterables at once

Starting with a vowel

This is the get_vowel_names exercise in loops.py.

Edit the function get_vowel_names so that it accepts a list of names and returns a new list containing all names that start with a vowel. It should work like this:

>>> from loops import get_vowel_names
>>> names = ["Alice", "Bob", "Christy", "Jules"]
>>> get_vowel_names(names)
['Alice']
>>> names = ["Scott", "Arthur", "Jan", "elizabeth"]
>>> get_vowel_names(names)
['Arthur', 'elizabeth']

Flatten a Matrix

This is the flatten exercise in loops.py.

Edit the function flatten, that will take a matrix (a list of lists) and return a flattened version of the matrix.

>>> from loops import flatten
>>> matrix = [[row * 3 + incr for incr in range(1, 4)] for row in range(4)]
>>> matrix
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
>>> flatten(matrix)
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

Matrix From String

This is the matrix_from_string exercise that we’ve seen before in loops.py. The function accepts a string and returns a list of lists of integers (found in the string).

Refactor the function to use a list comprehension.

Example:

>>> from loops import matrix_from_string
>>> matrix_from_string("1 2\n10 20")
[[1, 2], [10, 20]]

Power List By Index

This is the power_list exercise in ranges.py.

Edit the function, power_list, so that it accepts a list of numbers and returns a new list that contains each number raised to the i-th power where i is the index of that number in the given list. For example:

>>> from ranges import power_list
>>> power_list([3, 2, 5])
[1, 2, 25]
>>> numbers = [78, 700, 82, 16, 2, 3, 9.5]
>>> power_list(numbers)
[1, 700, 6724, 4096, 16, 243, 735091.890625]

Matrix Addition

This is the matrix_add exercise we saw earlier in ranges.py.

The function matrix_add accepts two matrices (lists of lists of numbers) and returns one matrix that includes each corresponding number in the two lists added together.

You should assume the lists of lists provided will always be the same size/shape.

Re-write the function to use list comprehensions.

>>> from ranges import matrix_add
>>> m1 = [[1, 2], [3, 4]]
>>> m2 = [[5, 6], [7, 8]]
>>> matrix_add(m1, m2)
[[6, 8], [10, 12]]
>>> m1 = [[1, 2, 3], [0, 4, 2]]
>>> m2 = [[4, 2, 1], [5, 7, 0]]
>>> matrix_add(m1, m2)
[[5, 4, 4], [5, 11, 2]]

Identity Matrix

This is the identity exercise in ranges.py.

Edit the identity function so that it takes as input a number size for the size of the matrix and returns an identity matrix of size x size elements.

An identity matrix is a square matrix with ones on the main diagonal and zeros elsewhere. A 3 by 3 identity matrix looks like:

identity_matrix = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]

The identity function should work like this:

>>> from ranges import identity
>>> identity(3)
[[1, 0, 0], [0, 1, 0], [0, 0, 1]]
>>> identity(4)
[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
>>> identity(2)
[[1, 0], [0, 1]]

X marks the spot

This is the words_containing exercise in loops.py.

Edit the words_containing function so that it accepts a list of words and a letter and returns a list of all words in the string that contain the given letter.

Example:

>>> from loops import words_containing
>>> words_containing(['My', 'life', 'is', 'my', 'message'], 'y')
['My', 'my']
>>> words_containing(['My', 'life', 'is', 'my', 'message'], 'i')
['life', 'is']
>>> words_containing(['My', 'life', 'is', 'my', 'message'], 'm')
['My', 'my', 'message']

Pythagorean Triples

This is the triples exercise in ranges.py.

Edit the triples function so that it takes a number and returns a list of tuples of 3 integers where each tuple is a Pythagorean triple, and the integers are all less then the input number.

A Pythagorean triple is a group of 3 integers a, b, and c, such that they satisfy the formula a**2 + b**2 = c**2

>>> from ranges import triples
>>> triples(15)
[(3, 4, 5), (5, 12, 13), (6, 8, 10)]
>>> triples(30)
[(3, 4, 5), (5, 12, 13), (6, 8, 10), (7, 24, 25), (8, 15, 17), (9, 12, 15), (10, 24, 26), (12, 16, 20), (15, 20, 25), (20, 21, 29)]