List Comprehension

A Simple Introduction to List Comprehension in Python

List comprehension is a compact way of creating lists. The simple formula is [ expression + context ].

  • Expression: What to do with each list element?
  • Context: What list elements to select? It consists of an arbitrary number of for and if statements.

The example [x for x in range(3)] creates the list [0, 1, 2].

In this article, you will learn everything you need to know about list comprehension in Python. Being hated by newbies, experienced Python coders can’t live without this awesome Python feature.

What is list comprehension?

List comprehension is a concise way of creating lists. Say you want to filter out all customers from your database who earn more than $1,000,000. This is what a newbie not knowing list comprehension would do:

# (name, $-income)
customers = [("John", 240000),
             ("Alice", 120000),
             ("Ann", 1100000),
             ("Zach", 44000)]

# your high-value customers earning >$1M
whales = []
for customer, income in customers:
    if income>1000000:
        whales.append(customer)
print(whales)
# ['Ann']

This snippet needs four lines just to create a list of high-value customers (whales)!

If you do that in your public Python code base, be prepared to get busted for “not writing Pythonic code”. 😉

Instead, a much better way of doing the same thing is to use list comprehension:

whales = [x for x,y in customers if y>1000000]
print(whales)
# ['Ann']

Beautiful, isn’t it?

List comprehension is dead simple when you know the formula I will show you in a moment. So why are people confused about how to use list comprehension? Because they never looked up the most important statement on list comprehension in the Python documentation. It’s this:

“A list comprehension consists of brackets containing an expression followed by a for clause, then zero or more for or if clauses. The result will be a new list resulting from evaluating the expression in the context of the for and if clauses which follow it.”

Official Python Documentation

In other words, here is the formula for list comprehension. That’s the one thing you should take home from this article.

Formula: List comprehension consists of two parts.

‘[‘ + expression + context + ‘]’

The first part is the expression. In the example above it was the variable x. But you can also use a more complex expression such as x.upper(). Use any variable in your expression that you have defined in the context within a loop statement. See this example:

whales = [x.upper() for x,y in customers if y>1000000]
print(whales)
# ['ANN']

The second part is the context. The context consists of an arbitrary number of for and if clauses. The single goal of the context is to define (or restrict) the sequence of elements on which we want to apply the expression. That’s why you sometimes see complex restrictions such as this:

small_fishes = [x + str(y) for x,y in customers if y<1000000 if x!='John']
# (John is not a small fish...)
print(small_fishes)
# ['Alice120000', 'Zach44000']

That’s about it!

List comprehension is easy once you have invested one or two coffee breaks into your thorough understanding. Consider this done! If you wonder why I discretize time into short 5-minute slots called “coffee breaks”, read my book “Coffee Break Python” ;).

To sum up, remember this one formula from this article: list comprehension = ‘[‘ + expression + context + ‘]’.

How does nested list comprehension work in Python?

After publishing the first version of this article, many readers asked me to write a follow-up article of nested list comprehension in Python. There are two interpretations of nested list comprehension:

  • Coming from a computer science background, I was assuming that “nested list comprehension” refers to the creation of a list of lists. In other words: How to create a nested list with list comprehension?
  • But after a bit of research, I learned that there is a second interpretation of nested list comprehension: How to use a nested for loop in the list comprehension?

How to create a nested list with list comprehension?

It is possible to create a nested list with list comprehension in Python. What is a nested list? It’s a list of lists. Here is an example:

## Nested List Comprehension
lst = [[x for x in range(5)] for y in range(3)]
print(lst)
# [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]

As you can see, we create a list with three elements. Each list element is a list by itself.

Everything becomes clear when we go back to our magic formula of list comprehension: [ expression + context]. The expression part generates a new list consisting of 5 integers. The context part repeats this three times. Hence, each of the three nested lists has five elements.

If you are an advanced programmer, you may ask whether there is some aliasing going on here. Aliasing in this context means that the three list elements point to the same list [0, 1, 2, 3, 4]. This is not the case because each expression is evaluated separately, a new list is created for each of the three context executions. This is nicely demonstrated in this code snippet:

l[0].append(5)
print(l)
# [[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]
# ... and not [[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]]

How to use a nested for loop in the list comprehension?

To be frank, the latter one is super-simple stuff. Do you remember the formula of list comprehension (= ‘[‘ + expression + context + ‘]’)?

The context is an arbitrary complex restriction construct of for loops and if restrictions with the goal of specifying the data items on which the expression should be applied.

In the expression, you can use any variable you define within a for loop in the context. Let’s have a look at an example.

Suppose you want to use list comprehension to make this code more concise (for example, you want to find all possible pairs of users in your social network application):

# BEFORE
users = ["John", "Alice", "Ann", "Zach"]
pairs = []
for x in users:
    for y in users:
        if x != y:
            pairs.append((x,y))
print(pairs)
#[('John', 'Alice'), ('John', 'Ann'), ('John', 'Zach'), ('Alice', 'John'), ('Alice', 'Ann'), ('Alice', 'Zach'), ('Ann', 'John'), ('Ann', 'Alice'), ('Ann', 'Zach'), ('Zach', 'John'), ('Zach', 'Alice'), ('Zach', 'Ann')]

Now, this code is a mess! How can we fix it? Simply use nested list comprehension!

# AFTER
pairs = [(x,y) for x in users for y in users if x!=y]
print(pairs)
# [('John', 'Alice'), ('John', 'Ann'), ('John', 'Zach'), ('Alice', 'John'), ('Alice', 'Ann'), ('Alice', 'Zach'), ('Ann', 'John'), ('Ann', 'Alice'), ('Ann', 'Zach'), ('Zach', 'John'), ('Zach', 'Alice'), ('Zach', 'Ann')]

As you can see, we are doing exactly the same thing as with un-nested list comprehension. The only difference is to write the two for loops and the if statement in a single line within the list notation [].

Is it possible to use the ‘else’ statement in Python list comprehension?

Yes, you can use the else statement in list comprehension – but only in the expression and not in the context part of list comprehension, assuming [ expression + context ].

lst = [x if x%2==0 else x-1 for x in range(10)]
# [0, 0, 2, 2, 4, 4, 6, 6, 8, 8]

To highlight the expression and the context parts, let’s use some parenthesis around the expression:

lst = [(x if x%2==0 else x-1) for x in range(10)]
# [0, 0, 2, 2, 4, 4, 6, 6, 8, 8]

However, you can not use the else statement within the context part. When I executed this code snippet, I was surprised that it throws an error:

# WRONG (INVALID SYNTAX)
lst = [x for x in range(10) if x%2==0 else x-1]
print(lst)

Why does this fail? The reason is the misplaced else statement. Without the else statement, the code executes properly. I can not really think of any convincing argument supporting this design decision of not allowing an else branch within the context. But as there is an easy work-around (moving it to the expression part as shown above), you won’t find it too difficult to avoid this situation.

What is the difference between list comprehension and set comprehension in Python?

There are two differences between list comprehension and set comprehension. https://blog.finxter.com/brackets-a-simple-introduction-to-set-comprehension-in-python/.

  • Braces vs. Brackets: Do you want to generate a set? Use curly braces {}. Do you want to generate a list? Use angle brackets [].
  • The data type of the result: list comprehension generates a list and set comprehension generates a set.

But syntactically, list comprehension and set comprehension are identical.

Where to go from here?

This article is originally published in my Python Email Course: “Increase Your Value to the Marketplace – in your coffee break”. For my subscribers, I regularly publish educative emails about the most important Python topics. Register and join my community of thousands of ambitious coders. I guarantee, you will love it! (It’s free.)