Skip to content

#30: List Comprehension

If you want to create a list from the values of another list, you most likely will write a for loop. That works and is how you would solve this exercise in any other language. However, there is a more pythonic way to do it called list comprehension.

List comprehension?

Python allows you to write elegant and easy to read code that almost looks like English. When you work with lists the feature called list comprehension can help you to replace a for loop with a single line of code using this syntax:

[ (expression) for (value) in (collection) if (condition) ]

Calculation the square of the even numbers between 1 and 5 can look like this example:

1
2
>>> [x*x for x in [1,2,3,4,5] if x % 2 == 0]
[4, 16]

As I saw this the first time I was intrigued and baffled at the same time. It looked so easy to read but I had absolutely no clue how to write it. Thanks to a great video as part of the Python Tricks Digital Toolkit by Dan Bader things started to make sense. Put all this information for a minute away and let us look at for loops and how we can rewrite them.

Square some values from a list

If you want to square the numbers in a list, your approach with a for loop may look like this:

1
2
3
4
5
6
squares = []
for x in range(10):
    squares.append(x*x)

print(squares)
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

As explained in an older post, the range function here gives you the numbers between 0 up to 10 (but not the 10 itself). You could write the numbers explicitly, but the range() function shows nicely that you can use any collection and not just lists.

With list comprehension, you can write this code for the same result:

1
2
3
squares = [x*x for x in range(10)]
print(squares)
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

In this example x*x is the expression, the x in for x in is the value and range(10) is the collection.

Add a filter

If we only want to get the squares of even numbers, we can add an if clause and use the modulo operator:

1
2
3
4
5
6
7
squares = []
for x in range(10):
    if(x % 2 == 0):
        squares.append(x*x)

print(squares)
[0, 4, 16, 36, 64]

With list comprehension, we can use the same condition at the end:

1
2
3
squares = [x*x for x in range(10) if x % 2 == 0]
print(squares)
[0, 4, 16, 36, 64]

This example shows how much code you can replace with this feature. Everything you need fits into one single line. As always, do not overuse such features. The more work you do inside your loop, the more likely it is that a loop is a better solution.

All together in one image

To better remember the different parts and where they go, I created this image:

The same code needs to go in a different place, but otherwise it works the same

Conclusion

List comprehensions are a powerful tool in Python that took me a while to understand. I hope this explanation and the graphic help you to write your own list comprehensions.