Comprehensions

For Loops

Let’s take a look at Python’s for loops.

Python’s for loops are not the same as for loops in Java, C++, Javascript, or many other programming languages. Python has a ‘for in’ loop. What we call a for loop in Python—our for in loop—is most often called a ‘for each’ loop in other programming languages.

A very simple for loop in Python might look like this:

numbers = [1, 2, 3, 5, 7]
for number in numbers:
    print(number)

In this for loop we are looping over every item in our list and printing those items out. There’s no index, there’s no looking for the length of the list or index incrementing. Python’s for loops magically do all that work for us under the hood.

List Comprehensions

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

>>> my_favorite_numbers = [2, 1, 3, 4, 7, 11, 18]
>>> doubled_numbers = []
>>> for n in my_favorite_numbers:
...     doubled_numbers.append(n * 2)
...
>>> doubled_numbers
[4, 2, 6, 8, 14, 22, 36]

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
[4, 2, 6, 8, 14, 22, 36]

This is called a list comprehension. List comprehensions provide convenient shorthand for creating lists from other lists or iterables. List comprehensions can be written in one line as shown above, or they can be broken up over multiple lines for readability:

>>> doubled_numbers = [
...     n * 2
...     for n in my_favorite_numbers
... ]
>>> doubled_numbers
[4, 2, 6, 8, 14, 22, 36]

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]

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]

List Comprehension Exercises

These exercises are all in the lists.py file in the exercises directory. Edit the file to add the functions or fix the error(s) in the existing function(s).

To run the test: from the exercises folder, type python test.py <function_name>, like this:

$ python test.py get_vowel_names

Tip

Start with a for loop and then copy-paste your way into a list comprehension.

Starting with a vowel

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:

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

Flatten a Matrix

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

Edit the function matrix_from_string so it accepts a string and returns a list of lists of integers (found in the string).

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

Power List By Index

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 lists 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

Edit the function matrix_add so it 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.

>>> 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

Edit the function identity 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. It should work like this:

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

[[1, 0, 0], [0, 1, 0], [0, 0, 1]]
>>> from lists 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]]

Pythagorean Triples

Edit the function triples 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 lists 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)]
Write more Pythonic code

I send out 1 Python exercise every week through a Python skill-building service called Python Morsels.

If you'd like to improve your Python skills every week, sign up!

You can find the Privacy Policy here.
reCAPTCHA protected (Google Privacy Policy & TOS)