List Comprehensions in Python 3 for Beginners

So you’re learning to code in Python, and you’re starting to understand how to work with different data types, variables, conditional statements, for loops, while loops, and maybe you’re feeling pretty good about it.

And then you see something like this:

[x+1 if x >= 5 else x+2 for x in l]

What in the world is this all about?? Is it a loop? An if statement? And how is this all in one line?

This is an example of list comprehension in Python. Contrary to what the name might suggest, being able to comprehend code like this is probably one of the biggest disadvantages to using this tool, however, it can be very powerful to have in your arsenal. For one thing, something that might take 5 or 6 lines of code can be reduced to one concise line. Also, list comprehensions are more computationally efficient under the hood, and once you get over the hump of understanding how they work, they can make your code simpler and more elegant.

Here are the basics of list comprehensions, and some ways in which they can be used to make your life easier when coding in Python.

What is List Comprehension?

The basic idea of list comprehension is to take a group of elements, modify the elements in some way, and return them in a list. Simple right? Well, there's a little more to it than that…

The general idea of a simple list comprehension syntactically looks like this:

[do_something for each_thing in a_group_of_things]

In other words, the do_something part should represent what you want for your output. The each_thing part represents the iterating variable or each of the original elements as you loop through them. The a_group_of_things part represents the input, or what we are iterating over. It can be anything that contains a group of elements. Most commonly this would be a list, but it can be anything from a tuple, a dictionary, a pandas series, a set, or even a string (a group of characters). If you think about everything in the above example after the do_something, it looks and behaves pretty much the exact same way a simple for loop would. So this is analogous to:

for each_thing in a_group_of_things:
    do_something

A crucial thing to understand about list comprehensions, however, is that no matter what the input, or what functions you perform on the input, the output of a list comprehension will always be a list, as the name implies. The order of the syntax can be a little confusing at first, but once you understand it, it can make a whole slew of tasks much easier.

How to Use it in Practice

Now that we have gone over the basic concept of a list comprehension, let’s take a really simple example to more clearly understand what is going on. Say we have the following list of numbers:

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

I want to create a new list from this which adds 1 to each of the numbers in the original list. There are several ways we can do this. We could, for example, tackle this using a for loop. This is what that would look like:

new_numbers = []
for number in numbers:
    number += 1
    new_numbers.append(number)

The new_numbers list will now look like this:

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

This is fine, but we can also do this much more elegantly using a list comprehension:

new_numbers = [number + 1 for number in numbers]

That's it! We have effectively reduced four lines into one and made less work for ourselves and our computers. We don’t need to worry about setting up a for loop or creating an empty list and individually appending each element. Easy right?

Beyond just simple addition, we can perform much more complex tasks in a list comprehension if we want to. Here’s another example using the same numbers list from above:

new_numbers = [round(number / 3) for number in numbers]

In this example, we are taking each number from our original numbers list, dividing it by 3, and then rounding to the nearest integer. The result of new_numbers will now look like this:

[0, 0, 1, 1, 1, 2, 2, 2, 3, 3]

What if I wanted to make the numbers into letters “a” through “j” using a list comprehension. I could make a dictionary that has the numbers as keys and the corresponding letters as values. Then I can use that in a list comprehension to generate a list of letters, like so:

numbers_to_letters = {
    0: 'a',
    1: 'b',
    2: 'c',
    3: 'd',
    4: 'e',
    5: 'f',
    6: 'g',
    7: 'h',
    8: 'i',
    9: 'j'
}
letters = [numbers_to_letters[num] for num in numbers_to_letters]

The letters list will now output to:

['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']

We can even call functions within a list comprehension. This is especially useful because we can make a function as complex as we like and then easily implement it on a group of elements with one simple line. Here’s an example:

def remove_vowels(word):
    vowels = ['a', 'e', 'i', 'o', 'u']
    for letter in word:
        if letter in vowels:
            word = word.replace(letter, '')
    return word

words = ['car', 'house', 'plant', 'fisherman', 'picnic', 'rodeo']
new_words = [remove_vowels(word) for word in words]

Here we are creating a function to remove the vowels from a string and then using a list comprehension to run that function on each of the words in the words list.

new_words will now output to:

['cr', 'hs', 'plnt', 'fshrmn', 'pcnc', 'rd']

This can be done will any number of functions or methods, which is part of what makes list comprehension a very pythonic way of performing complex operations on a large group of items.

Conditionals Within List Comprehensions

List comprehensions also allow us to use if and else statements, giving them even more specific functionality. This is pretty awesome, although once again the syntax can be a little tricky. Say I want to use our original example, but this time only change the odd numbers and leave the rest of the numbers as they are. We can do this by simply placing if and else statements inside of a list comprehension:

numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
new_numbers = [number + 1 if number % 2 == 1 else number for number in numbers]

Now we have the following for new_numbers:

[0, 2, 2, 4, 4, 6, 6, 8, 8, 10]

The order of the if else and for may look a little strange at first. For conditional logic in a list comprehension like the example above, you want to think about it like this:

[result_if_condition_met if condition else result_if_condition_not_met for element in group_of_elements]

An important note here is the without the else statement, anything that does not meet the condition will not be returned at all. The list comprehension will remove that item from the output list entirely. There are times when we might actually want to do this. This technique is known as filtering, which is using list comprehension to only return certain items in a group based on a specific condition or a set of conditions. Unfortunately, Python doesn’t make it easy for us when it comes to syntax. This is one of the most confusing things to get straight about list comprehensions. When you’re not using an else statement, the if statement needs to be moved to the end of the list comprehension, like so:

[result for element in group_of_elements if condition]

Here’s a simple filtering example for our above numbers list:

new_numbers = [number for number in numbers if number % 2 == 1]

Now new_numbers will only contain the odd numbers:

[1, 3, 5, 7, 9]

Here we are just passing each original number to the new list without making any changes to them, as long as they are odd numbers. However, we could have also modified the numbers as well as filtering them, for example:

new_numbers = [number + 1 for number in numbers if number % 2 == 1]

Output:

[2, 4, 6, 8, 10]

You can even filter based on multiple conditions, or in other words, nest one condition into another. For example, the following code:

new_numbers = [number for number in numbers if number % 2 == 1 if number < 5]

Yields this:

[1, 3]

This is going through each item in numbers filtering out only the odds, and from there, only returning numbers that are less than 5.

Now for the really fun part… multiple if and else statements in a single list comprehension. Yes, we actually can do that too!

This is conceptually a lot like using if, elif and else statements in a for loop, however in list comprehensions we don’t use elif and the syntax is of course very different. Let's take a look at the code below:

numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
new_numbers = []
for number in numbers:
    if number < 5:
        new_numbers.append(number)
    elif number == 7:
        new_numbers.append('seven')
    else:
        new_numbers.append('big')

This code will return the following list for new_numbers:

[0, 1, 2, 3, 4, 'big', 'big', 'seven', 'big', 'big']

Now here’s the equivalent of that in a list comprehension:

new_numbers = [number if number < 5 else 
               'seven' if number == 7 else 
               'big' for number in numbers]

This is much more concise and pythonic code and returns the exact same result. Rather than using an elif, we do this using if and else conditionals in this particular order. The second if acts just like an elif would in the for loop version of this. Here’s a more generalized way of thinking about multiple conditionals in list comprehension:

[result_1 if condition_1 else
 result_2 if condition_2 else
 result_3 if condition_3 else
 result_if_no_conditions_met for element in group_of_elements]

This can be generalized for any number of conditions which makes for a very useful way to incorporate complex conditional logic without having to write many lines of code. The syntax may be wonky, but once you get over that, the possibilities are endless.

Nested List Comprehensions

One more really cool thing you should know about list comprehensions is that you can actually nest one inside another. For example, if you have a list of lists, you could modify them using nested list comprehensions. This would look something like this:

[[do_something for item in inner_list] for inner_list in outer_list]

The outer set of brackets is iterating over the outer list and the inner set of brackets deals with each of the inner lists. This can be simplified even further as such:

[[do a list comp] for inner_list in outer_list]

Remember the example from before of removing all the vowels from a list of strings? Well, now that we have discovered nested list comprehensions, check out the difference we can make in solving this problem.

Long version:

words = ['car', 'house', 'plant', 'fisherman', 'picnic', 'rodeo']
new_words = []
for word in words:
    vowels = ['a', 'e', 'i', 'o', 'u']
    for letter in word:
        if letter in vowels:
            word = word.replace(letter, '')
    new_words.append(word)

Better:

words = ['car', 'house', 'plant', 'fisherman', 'picnic', 'rodeo']
vowels = ['a', 'e', 'i', 'o', 'u']
new_words = [''.join([letter for letter in word if letter not in vowels]) for word in words]

Output:

['cr', 'hs', 'plnt', 'fshrmn', 'pcnc', 'rd']

Both of these code blocks produce the exact same result. The first one uses nested for loops and the second one uses nested list comprehensions. As you can see, the second one is much more concise, and it’s less work on your computer too. This example in particular really demonstrates the power of list comprehension and shows why it’s so widely used in Python. We were able to iterate through each word in a list, remove all the vowels from them, and return them in a new list, all using just one line!

There are many other things that can be accomplished with list comprehensions, but hopefully, this will serve as a foundation for you to continue exploring the various applications for this nifty Python tool. If you can master list comprehension, it will really take your code to the next level.

Good Luck!