Python List Comprehension: Is It Easy to Comprehend?!?

Have you ever heard about the list comprehension in Python? It simplifies the way you work with lists and it makes your code more concise.

The list comprehension is a Python construct that reduces the lines of code required to generate a new list or to filter an existing list. A list comprehension is enclosed within square brackets and it’s made of an expression, one or more for loops and an optional condition to filter the list generated.

We will first define the list comprehension and then go through a series of examples that will make this part of your coding knowledge.

Let’s discover list comprehensions together!

What Does a List Comprehension Do?

Python list comprehensions allow to create a brand new list or to generate a list by filtering or mapping an existing list.

List comprehensions use the following syntax:

new_list = [expression(item) for item in iterable if condition]

Examples of iterables in Python are lists, tuples, sets and strings.

Given an iterable a list comprehension loops through the items in the iterable, applies the expression to each one of them and based on that generates a new list.

An optional condition can be also specified to filter the items in the iterable.

The result of a list comprehension is a new list that would require a lot more lines of code if you had to create it by using standard for loops and if statements.

Here is how the one line code above would look like without a list comprehension:

new_list = []

for item in iterable:
    if condition:
        new_list.append(expression(item))

A lot better on a single line!

A list comprehension is called this way because it’s a comprehensive or complete way to describe a sequence in Python.

How To Create a New List Using a List Comprehension

One of the main things you can do with a list comprehension is to create a new list.

For example, let’s see how to create a new list by using the range function inside a list comprehension.

>>> numbers = [x for x in range(10)]
>>> print(numbers)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

The expression of the list comprehension is very simple in this case, it’s just x.

Let’s update the expression to double the value of x:

>>> numbers = [2*x for x in range(10)]
>>> print(numbers)
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

As you can see we have created a list where every element is multiplied by 2.

The expression can be anything you want.

How to Add a Single Conditional Statement to a List Comprehension

Let’s start with the list comprehension in the previous section.

The next step to learn more about list comprehensions is to add a condition to it.

The syntax we will use is:

new_list = [expression(item) for item in iterable if condition]

Assume, for example, that we still want to generate the following list:

>>> numbers = [2*x for x in range(10)]

But this time we want to exclude numbers greater or equal to 5.

>>> numbers = [2*x for x in range(10) if x < 5]
>>> print(numbers)
[0, 2, 4, 6, 8]

We have filtered the elements in the new list by using a condition.

How to Add Two Conditions to a List Comprehension

To add two conditions to a list comprehension simply add both conditions at the end of the list comprehension one after the other (before closing the square bracket).

Update the previous list comprehension to only take into account numbers between 2 and 5 (2 and 5 excluded).

>>> numbers = [2*x for x in range(10) if x > 2 and x < 5]
>>> print(numbers)
[6, 8]

Makes sense?

How Do You Convert a For Loop into a List Comprehension?

Let’s start by defining a list of strings:

animals = ['tiger', 'lion', 'elephant']

I want to create a for loop that adds the letter ‘s’ at the end of each string to create a list of plurals.

>>> new_animals = []
>>> for animal in animals:
...     new_animals.append(animal + 's')
... 
>>> print(new_animals)
['tigers', 'lions', 'elephants']

Firstly we define a new empty list that we will use for the plural nouns. Then at each iteration of the for loop we use the append method to add a string to the new list.

This code works but is there a way to make it more concise?

Using a list comprehension we can simplify this code, here is how:

>>> new_animals = [animal + 's' for animal in animals]
>>> print(new_animals)
['tigers', 'lions', 'elephants']

That’s great!

With a single line of code, we have created a new list instead of using three lines of code like we have seen before.

When using a list comprehension we don’t have to create an empty list at the beginning.

In this list comprehension, the expression is animal + ‘s’ and it’s followed by a for loop that goes through the elements of the initial list one at a time and applies the expression to each one of them.

Can You Use Else in a List Comprehension?

In the previous example, we have used an if statement in a list comprehension.

But, can you also use an else statement to add multiple conditions to the list comprehension?

Let’s give it a try…

…start from the code below:

>>> numbers = [2*x for x in range(10) if x > 2 and x < 5]

To add an else condition we have to rearrange the order of the list comprehension elements.

We have to move the condition before the for keyword so we can return a value different than 2*x when the if condition is not met.

>>> numbers = [2*x if x > 2 and x < 5 else 3*x for x in range(10)]
>>> print(numbers)
[0, 3, 6, 6, 8, 15, 18, 21, 24, 27]

So, here is what happens in this code…

If the value of x is between 2 and 5 the list comprehension returns 2*x otherwise it returns 3*x.

For example, the number 1 is not between 2 and 5 and hence the result is 3*1 = 3.

Using Elif in a List Comprehension

It’s not possible to use the elif statement in a list comprehension but it’s possible to implement the same behavior by using multiple else statements.

Start with the following code:

>>> numbers = [2*x if x > 2 and x < 5 else 3*x for x in range(10)]

At the moment the condition is the following:

• if x > 2 and x < 5 => return 2*x
• else => return 3*x

I want to implement the following behaviour:

• if x > 2 and x < 5 => return 2*x
• else if x <=2 => return 3*x
• else => return 4*x

The conditional expression in a list comprehension is based on the ternary operator and we can use the following code to implement the behaviour we want.

>>> numbers = [2*x if x > 2 and x < 5 else 3*x if x <=2 else 4*x for x in range(10)]
>>> print(numbers)
[0, 3, 6, 6, 8, 20, 24, 28, 32, 36]

I know, it’s a long expression and at this point I would consider using a standard implementation instead of a list comprehension.

Writing code that works is not the only thing that matters…

It’s very important to write readable code because if code is not readable this can lead to bugs and can make managing existing code a nightmare.

This list comprehension can also be written in the following way (still without using elif):

numbers = []

for x in range(10):
    if x > 2 and x < 5:
        numbers.append(2*x)
    else:
        if x <=2:
            numbers.append(3*x)
        else:
            numbers.append(4*x)  

If you print the value of numbers you get the same result:

[0, 3, 6, 6, 8, 20, 24, 28, 32, 36]

This code is definitely more readable than the list comprehension and it can become even more readable if we use the elif statement:

numbers = []

for x in range(10):
    if x > 2 and x < 5:
        numbers.append(2*x)
    elif x <=2:
        numbers.append(3*x)
    else:
        numbers.append(4*x)

Execute this code on your computer and verify that the result is the same.

How to You Use the Break Statement in a List Comprehension

In a standard Python for loop you can stop the execution of the loop by using the break statement if a specific condition occurs.

How can you do the same with a list comprehension?

List comprehensions don’t support the break statement but it’s possible to use alternative approaches to simulate the behaviour of the break statement.

For example, let’s say we have a list of random numbers and we want to stop the execution of the list comprehension if a specific number is encountered.

Firstly let’s see how we can generate a list of random numbers without a list comprehension:

import random

random_numbers = []
while len(random_numbers) < 10:
    random_number = random.randint(1, 5)
    random_numbers.append(random_number)

We create an empty list and then append random numbers between 1 and 5 to it until the list of numbers has 10 elements.

Here is the output:

[1, 3, 5, 3, 2, 1, 3, 3, 4, 3]

Now let’s add a break statement to stop the execution of the while loop if the number 3 is encountered.

We will add the number 3 to the list before breaking from the loop. In this way, we can confirm the logic in our program works by seeing the number 3 as the last element of the new list of numbers.

import random

random_numbers = []
while len(random_numbers) < 10:
    random_number = random.randint(1, 5)
    random_numbers.append(random_number)

    if random_number == 3:
        break

The program works fine. You might have to run it few times in case the number 3 is not generated by random.randint.

[5, 3]

Now, let’s do the same with a list comprehension, start by generating the full list of 10 random numbers…

>>> random_numbers = [random.randint(1,5) for x in range(10)]
>>> print(random_numbers)
[2, 2, 4, 4, 4, 1, 3, 5, 2, 4]

Once again the list comprehension rocks! A single line replaces multiple lines of code.

And now, how do we stop the list comprehension if the number 3 is encountered?

One possible approach requires an external module: itertools. We will use the function itertools.takewhile().

First we generate the random_numbers list.

>>> import itertools
>>> random_numbers = [random.randint(1,5) for x in range(10)]
>>> print(random_numbers)
[2, 3, 5, 4, 5, 4, 2, 5, 3, 4]

Then we pass it to the function itertools.takewhile.

>>> print(itertools.takewhile(lambda number: number !=3, random_numbers))
<itertools.takewhile object at 0x7f88a81fe640>

The function itertools.takewhile takes as:

• first argument a lambda that defines the condition for the program to continue running.
• second argument the iterable.

It returns an itertools.takewhile object that we have to convert into a list to see the elements.

>>> print(list(itertools.takewhile(lambda number: number !=3, random_numbers)))
[2]

The code does what we want. At the same time, the behaviour is not exactly identical to the one using a break statement.

That’s because we first generate the full list of random numbers and then go through them until the number 3 is encountered.

Also in the second implementation, the number 3 is not included in the final list.

In this scenario, it’s definitely a lot easier to use the break statement than a convoluted list comprehension that requires itertools.takewhile and a lambda.

It’s a bit too much! 😀

Use a List Comprehension With Two or More Lists

One way to apply a list comprehension to two or more lists is to use it together with the zip() function.

>>> cities = ['Rome', 'Warsaw', 'London']
>>> countries = ['Italy', 'Poland', 'United Kingdom']
>>> [(city, country) for city, country in zip(cities, countries)]
[('Rome', 'Italy'), ('Warsaw', 'Poland'), ('London', 'United Kingdom')]

A list comprehension used with the zip function returns a list of tuples where the nth tuple contains the nth element of each list.

The same applies if we pass three lists to the list comprehension (and so on).

>>> cities = ['Rome', 'Warsaw', 'London']
>>> countries = ['Italy', 'Poland', 'United Kingdom']
>>> languages = ['Italian', 'Polish', 'English']
>>> [(city, country, language) for city, country, language in zip(cities, countries, languages)]
[('Rome', 'Italy', 'Italian'), ('Warsaw', 'Poland', 'Polish'), ('London', 'United Kingdom', 'English')]

Replace Map and Lambda with a List Comprehension

The map function applies a given function to the elements of an iterable.

For example, you can use the map function to double the value of every number in a list.

>>> numbers = [3, 6, 8, 23]
>>> print(map(lambda x: 2*x, numbers))
<map object at 0x7f88a820d2e0>
>>> print(list(map(lambda x: 2*x, numbers)))
[6, 12, 16, 46]

Notice that the first argument passed to the map function is a lambda function.

And here is how you can write this expression using a list comprehension.

>>> [2*x for x in numbers]
[6, 12, 16, 46]

Super easy!

Use a List Comprehension Instead of Filter and Lambda Functions

Using the filter function you can filter the elements of a list based on a given condition.

For example, let’s filter from the previous list of numbers those that are smaller than 10.

The condition is passed as the first argument to the filter function and it’s expressed as a lambda function.

>>> print(filter(lambda x: x<10, numbers))
<filter object at 0x7f88a8202340>
>>> print(list(filter(lambda x: x<10, numbers)))
[3, 6, 8]

And now write the same logic using a list comprehension.

>>> [x for x in numbers if x < 10]
[3, 6, 8]

Replace Reduce and Lambda with a List Comprehension

The reduce function applied to our list of numbers returns the overall sum based on the fact that we are using the following lambda function:

lambda a,b: a+b

Here is the result of the call to the reduce function:

>>> from functools import reduce
>>> numbers = [3, 6, 8, 23]
>>> print(reduce(lambda a,b: a+b, numbers))
40

If you want to understand how this works exactly have a look a this tutorial about the reduce function.

Now, we will convert it into a list comprehension. To obtain the same result we also have to use the sum() function.

>>> print(sum([number for number in numbers]))
40

How to Use Nested List Comprehensions

Nested list comprehensions can be useful when working with lists of lists.

For example, let’s say we want to write code that increases every number in a matrix by one.

This is our original matrix:

matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

Using for loops we would do the following:

for row in matrix:
    for index in range(len(row)):
        row[index] += 1

The updated matrix is:

[[2, 3, 4], [5, 6, 7], [8, 9, 10]]

How can we use a list comprehension instead of two nested loops?

We could try by simply translating the code above into a list comprehension.

>>> [[row[index] + 1 for index in range(len(row))] for row in matrix]
[[2, 3, 4], [5, 6, 7], [8, 9, 10]]

Notice that we have used one list comprehension inside another list comprehension. That’s why these are called nested list comprehensions.

Difference Between List Comprehension and Generator Expression

A Python construct that looks very similar to a list comprehension is the generator expression.

To convert a list comprehension into a generator expression replace the square brackets with parentheses.

Let’s see how this can be applied to the list of random numbers we have used before.

Remember to import the random module before running the following code otherwise, you will see a NameError exception.

List comprehension

>>> random_numbers = [random.randint(1,5) for x in range(10)]
>>> print(random_numbers)
[1, 4, 3, 5, 3, 4, 5, 4, 5, 4]
>>> print(type(random_numbers))
<class 'list'>

Generator expression

>>> random_numbers = (random.randint(1,5) for x in range(10))
>>> print(random_numbers)
<generator object <genexpr> at 0x7fccb814e3c0>
>>> print(type(random_numbers))
<class 'generator'>

As you can see when using a list comprehension we can print the full list of elements in the list generated.

The same doesn’t apply to the generator expression that just returns a generator object.

To get the next item from the generator object we have to use the Python next function:

>>> print(next(random_numbers))
3
>>> print(next(random_numbers))
2

The main difference between a list comprehension and a generator expression is in the way they store data in memory. The list comprehension returns immediately the full list of numbers. The generator expression creates a generator that returns one number at a time and hence optimizes memory usage.

For Loop vs List Comprehension: A Speed Comparison

In this last section of this tutorial, I want to perform a speed comparison between a for loop and a list comprehension when working on the same numbers.

Create a Python file called for_loop_vs_list_comprehension.py with the following code:

def get_numbers_using_for_loop():
    numbers = []

    for x in range(10):
        numbers.append(2*x)

    return numbers


def get_numbers_using_list_comprehension():
    numbers = [2*x for x in range(10)]
    return numbers

And confirm that both functions return the same result:

print(get_numbers_using_for_loop())
print(get_numbers_using_list_comprehension())

[output]
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

Then use the timeit module to measure the speed of both functions:

$ python -m timeit -s "from for_loop_vs_list_comprehension import get_numbers_using_for_loop" "get_numbers_using_for_loop()"
500000 loops, best of 5: 868 nsec per loop

$ python -m timeit -s "from for_loop_vs_list_comprehension import get_numbers_using_list_comprehension" "get_numbers_using_list_comprehension()"
500000 loops, best of 5: 731 nsec per loop

The implementation using the list comprehension is faster than the one that uses the for loop.

Conclusion

We have learned quite a lot about list comprehension in Python!

It’s great how a list comprehension can make your code a lot more concise and how it can also replace multiple Python constructs based on for loops, lambdas, and map / reduce / filter functions.

Are you ready to start using list comprehensions now?

If not go through this article again and practice, practice, practice 🙂