List Comprehensions in Python

A more elegant and concise way to create lists in python

Photo by Andrew Neel on Unsplash

Creating a List

Let’s say that we want to create a list in python from another iterable object or sequence, such as another list. For example, we have a list of numbers, and we want to create a new list that contains the cubes of the numbers from the first list. There are multiple ways to accomplish this task. Perhaps the most basic way is to use a for loop as follows:

In the above code, we use a for loop to loop over our num_list, and we add each number cubed to cube_list.

Well, it turns out that there’s an easier way to accomplish the same task, and that’s with using a list comprehension.

List Comprehensions

List comprehensions allow us to create lists from other sequences in a very concise way. We usually use list comprehensions to loop through another sequence, such as a list, and either add the elements that satisfy a certain condition, or add the result of an operation applied to each element in the sequence.

Writing a List Comprehension

A list comprehension is made up of brackets, which contain an expression, followed by a for loop, and zero or more for or if clauses. This list comprehension then creates a list that contains the expression evaluated in the context of the for and if clauses that follow it.

Let’s start with a basic list comprehension that contains only an expression and a for loop:

[ <expression> for <name> in <iterable or sequence> ]

For example, let’s see how we can create the above cube_list, or a list that contains the cubes of another list, using a list comprehension:

So what is going on in this line of code?

cube_list = [num**3 for num in num_list]

First, we notice that we have brackets that contain an expression, num**3 , followed by a for loop, for num in num_list . Within the for loop, num is the parameter name we give for the elements that we are looping over in our num_list , just like in the original for loop above. We are basically saying, take num (or the current element) from our num_list , cube it, and then add the result of that operation to our list, similar to cube_list.append( num**3 ). Thus we are adding the output of the expression num**3 to the list we are making as it iterates over the for loop.

Note: A list comprehension behaves very similarly to the built-inmap function in python.

Note on Expressions:

The expression in a list comprehension can include functions as well. For example, if we want to create a list that contains the corresponding lengths of a list of strings, we can do so with the following list comprehension:

list_of_strings = [‘hello’, ‘my’, ‘name’, ‘a’]len_of_strings = [len(word) for word in list_of_strings]print(len_of_strings) # output is [5,2,4,1]

Notice how we used the built-in len function as part of our expression within the list comprehension.

List Comprehension with Condition

We can also add a condition to our list comprehensions using an if statement after the for loop:

[ <expression> for <name> in <iterable or sequence> if <condition> ]

For example, let’s say that we want to make a new list from num_list that only contains the cubes of the odd numbers. Well, we accomplish that with the following list comprehension:

cubes_of_odds = [num**3 for num in num_list if num%2 != 0]print(cubes_of_odds) # output is [1,2,27,4,125]

Notice how we added the if statement at the end. So num**3 will only be added to our cubes_of_odds list if the current element or num is odd using the modulo operator. The modulo operator returns the remainder when num is divided by 2, which would equal zero if num is even.

What if we want multiple conditions?

Not only can we add an if statement, but we can add an else statement as well using the following format:

[ <expression> if <condition> else <expression> for <name> in <iterable or sequence> ]

For example, let’s say we want to loop through num_list, and make a new list with the cubes of the odd numbers and the squares of the even numbers. We can do that with the following code:

cubes_and_squares = [num**3 if num%2!=0 else num**2 for num in num_list]print(cubes_and_squares) # output is [1,4,27,16,125]

So as we loop through num_list , if num%2!=0 is True for that specific num or element, the num**3 is used as the expression for that specific element. If num%2!=0 is not True, the num**2 expression will be used for the element instead.

Nested List Comprehensions

The expression in a list comprehension can also be another list comprehension. For example, let’s say that we want to make the following matrix (or 2 dimensional array):

matrix = [[1, 2, 3, 4],
          [1, 2, 3, 4],
          [1, 2, 3, 4],
          [1, 2, 3, 4],
          [1, 2, 3, 4]]

Note that it consists of four rows, each row containing the numbers 1 through 4. In other words it is a list that contains four identical lists, each being [1,2,3,4].

We can make this matrix with a list comprehension in addition to an outer for loop as follows:

matrix = []for y in range(1,5):
    matrix.append([x for x in range(1,5)])

The for loop has four iterations, and in each iteration a list is created with the list comprehension: [x for x in range(1,5)] . This list comprehension creates a list of [1,2,3,4]. Thus we are creating a list that contains four lists of [1,2,3,4].

The above code is equivalent to the following nested list comprehension:

matrix = [[x for x in range(1,5)] for y in range(1,5)]

So we are executing the initial expression, which is a list comprehension, [x for x in range(1,5)], that creates a [1,2,3,4] list. This expression is executed with each iteration of the second for loop, each time creating a [1,2,3,4] list and adding it to the outer list.

Conclusion

In this tutorial, we learned how list comprehensions can be used to make lists from other lists or sequences. We saw how they can contain zero or more conditions. And lastly, we saw how to use a nested list comprehension to make a 2-dimensional list.