Python List Comprehension: Tutorial With Examples

There’s a concept called set-builder notation in mathematics, also called set comprehension. Inspired by this principle, Python offers list comprehensions. In fact, Python list comprehensions are one of the defining features of the language. It allows us to create concise, readable code that outperforms the uglier alternatives like for loops or using map().

We’ll first look at the most well-known type: list comprehensions. Once we’ve got a good grasp of how they work, you’ll also learn about set comprehensions and dictionary comprehensions.

What are list comprehensions?

A Python list comprehension is a language construct. It’s used to create a Python list based on an existing list. Sounds a little vague, but after a few examples, that ‘ah-ha!’ moment will follow, trust me.

The basic syntax of a list comprehension is:

[ <expression> for item in list if <conditional> ]

The ‘if’-part is optional, as you’ll discover soon. However, we do need a list to start from. Or, more specifically, anything that can be iterated. We’ll use Python’s range() function, which is a special type of iterator called a generator: it generates a new number on each iteration.

Examples of list comprehensions

Enough theory, let’s look at the most basic example, and I encourage you to fire up a Python REPL to try this yourself:

>>> [x for x in range(1, 5)]
[1, 2, 3, 4]

Some observations:

• The expression part is just x
• Instead of a list, we use the range() function. We can use [1, 2, 3, 4] too, but using range() is more efficient and requires less typing for longer ranges.

The result is a list of elements, obtained from range(). Not very useful, but we did create our first Python list comprehension. We could just as well use:

>>> list(range(1,5))
[1, 2, 3, 4]

So let’s throw in that if-statement to make it more useful:

>>> [x for x in range(1,10) if x % 2 == 0]
[2, 4, 6, 8]

The if-part acts like a filter. If the condition after the if resolves to True, the item is included. If it resolves to False, it’s omitted. This way, we can get only the even numbers using a list comprehension.

So far, our expression (the x) has been really simple. Just to make this absolutely clear though, expression can be anything that is valid Python code and is considered an expression. Example:

>>> [x + 4 for x in [10, 20]]
[14, 24]

This expression adds four to x, which is still quite simple. But we could also have done something more complicated, like calling a function with x as the argument:

def some_function(a):
    return (a + 5) / 2
    
m = [some_function(x) for x in range(8)]
print(m)
# [2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0]

More advanced examples

You mastered the basics; congrats! Let’s continue with some more advanced examples.

Nested list comprehension

If expression can be any valid Python expression, it can also be another list comprehension. This can be useful when you want to create a matrix:

>>> m = [[j for j in range(3)] for i in range(4)]
>>> m
[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]

Or, if you want to flatten the previous matrix:

>>> [value for sublist in m for value in sublist]
[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]

The same, but with some more whitespace to make this clearer:

>>> [value
     for sublist in m
     for value in sublist]
[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]

The first part loops over the matrix m. The second part loops over the elements in each vector.

Alternatives to list comprehensions

The Python language could do without comprehensions; it would just not look as beautiful. Using functional programming functions like map() and reduce() can do everything a list comprehension can.

Another alternative is using for-loops. If you’re coming from a C-style language, like Java, you’ll be tempted to use for loops. Although it’s not the end of the world, you should know that list comprehensions are more performant and considered a better coding style.

Other comprehensions

If there are list comprehensions, why not create dictionary comprehension as well? Or what about set comprehensions? As you might expect, both exist.

Set comprehensions

The syntax for a Python set comprehension is not much different. We just use curly braces instead of square brackets:

{ <expression> for item in set if <conditional> }

For example:

>>> {s for s in range(1,5) if s % 2}
{1, 3}

Dictionary comprehensions

A dictionary requires a key and a value. Otherwise, it’s the same trick again:

>>> {x: x**2 for x in (2, 4, 6)}
{2: 4, 4: 16, 6: 36}

The only difference is that we define both the key and value in the expression part.