List comprehension is an easy way of creating lists within a single line of code. Suppose you have a list of integers and you want to create a new list in which the elements are square of the corresponding elements of the first list. You can do this in a single line of code using list comprehension.
We will look at how to create lists using list comprehension with examples and will also look at the cases where it should be avoided.
Using List Comprehension in Python
Let’s take an example in which we will create a list of the cubes of numbers from 1 to 10. From what we have learned so far, we can easily use a for loop to iterate range() and calculate cubes and store them in a list.
The following example prints the cube of all numbers from 1 to 10 (included) using a for loop.
cubes = []
for i in range(1, 11):
cubes.append(i**3)
print(cubes)
We created an empty list cubes to store the cubes. Then we iterate over range(1, 11) using a for loop. In each iteration, the cube of the number is calculated and the result is appended to the list cubes.
We can do the same in a single line of code using list comprehension as shown below.
cubes = [i**3 for i in range(1, 11)]
print(cubes)
If you closely look at [i**3 for i in range(1, 11)], you will find a for loop definition - for i in range(1, 11).
Let’s start with the for loop - for i in range(1, 11). It is a simple for loop through which we are iterating over range(1, 11).
Now in each iteration, i**3 is evaluated.
So in the first iteration, i is 1 and hence i**3 is also 1.
In the second iteration, i is 2 and hence i**3 is 8. This goes for all the ten iterations.
The interesting part is that the values we get after evaluation of i**3 in each iteration are added in a list and that list is returned at the end of the loop. The results of the evaluation of the expression in different iterations constitute the different elements of the list.
Thus, this method of creating the list cubes is called list comprehension.
Summing up, the expression of a list comprehension is enclosed within brackets [ ].
Inside [ ], we use a for loop which has a variable (i in the above example) for iteration. Before the for loop we also have an expression (i**3 in the above example) which is evaluated in each iteration and the result of this evaluation is added to a list which we are going to get at the end of the loop.
We have created a list using list comprehension. Now let’s look at the general syntax of list comprehension.
Python List Comprehension Syntax
new_list = [expression for_loop]
The list returned by list comprehension method is enclosed within brackets [ ]. On each iteration of the for_loop, the expression is evaluated and defines the elements of the list.
Comparing this syntax to the last example, i**3 is expression and for i in range(1, 11) is for_loop.
Python List Comprehension Examples
Let’s look at a few more examples.
The following example creates a list mylist of numbers from 1 to 5 (included) using a for loop.
mylist = []
for i in range(1, 6):
mylist.append(i)
print(mylist)
We can create the same list using list comprehension as shown below.
mylist = [i for i in range(1, 6)]
print(mylist)
In for i in range(1, 6), the variable i iterates over range(1, 6).
In each iteration, the expression i is evaluated. In the first iteration, i is 1. In the second iteration, i is 2, and this goes on till the fifth iteration.
The values of i in the five iterations constitutes the five elements of the list mylist.
Comparing this code to the syntax, i is expression and for i in range(1, 6) is for_loop.
Now let’s create a list having five items, each equal to 2, using list comprehension.
mylist = [2 for i in range(1, 6)]
print(mylist)
In for i in range(1, 6), the variable i iterates over range(1, 6). Therefore there are five iterations. In each iteration, the value 2 is appended to the list. Therefore, the list mylist has five elements, each equal to 2.
Comparing this code to the syntax, 2 is expression and for i in range(1, 6) is for_loop.
Using List Comprehension with Conditional Statements in Python
We can also use some condition with list comprehension. For example, if we want to create a list of all even numbers from 1 to 100, we can add a condition i%2==0 and only the elements passing this condition can be included in our list.
Let’s see how to implement list comprehension when an if statement is used in the body of a for loop.
The following example creates a list containing all even numbers from 1 to 10 (included).
even = []
for num in range(1, 11):
if num % 2 == 0:
even.append(num)
print(even)
We iterate over range(1, 11). In each iteration, it is checked if the number is divisible by 2. If it is, then it is appended to the list even.
Let’s create this list using list comprehension.
even = [num for num in range(1, 11) if num % 2 == 0]
print(even)
Here in the for loop for num in range(1, 11), the variable num iterates over range(1, 11).
In each iteration, the if condition num % 2 == 0 checks if num is even. If the condition is True (which means if num is even), then num is appended to the list, otherwise it is not appended.
On comparing this example with the example using for loop, you will understand the match.
Now let’s look at the syntax of list comprehension when some if condition is used.
Syntax
new_list = [expression for_loop if condition]
On each iteration of the for_loop, the if condition is checked. If the condition is True, then only the expression is evaluated and appended to the list.
Comparing this syntax to the last example, num is expression, for i in range(1, 11) is for_loop and and if num % 2 == 0 is if condition.
Python If...else with List Comprehension
Using if...else while creating lists using list comprehension has a slightly different syntax than using just if. Let’s look at the two syntaxes.
Syntax of List Comprehension including if statement
new_list = [expression for_loop if condition]
We have already seen two examples where if statements were used.
Syntax of List Comprehension including if..else statement
new_list = [expression1 if condition else expression2 for_loop]
Here, if the condition of if is True, then expression1 is evaluated, else expression2 is evaluated. Rest everything is the same.
Let’s see how if..else is implemented in list comprehension.
The following example stores the square of all even numbers and the cube of all odd numbers from 1 to 10 (included) in a list using a for loop.
result = []
for num in range(1, 11):
if num % 2 == 0:
result.append(num ** 2)
else:
result.append(num ** 3)
print(result)
In each iteration, we are checking if the number is even. If it is even, we are appending the square of the number, otherwise we are appending the cube of the number to the list.
This can be transformed using list comprehension as follows.
result = [num ** 2 if num % 2 == 0 else num ** 3 for num in range(1, 11)]
print(result)
In the for loop for num in range(1, 11), the variable num iterates over range(1, 11).
In each iteration, the if condition num % 2 == 0 checks if num is divisible by 2. If this condition is True, then the expression num ** 2 is evaluated and appended to the list. Otherwise, the expression num ** 3 is evaluated and appended to the list.
Comparing this code to the syntax, num ** 2 is expression1, num ** 3 is expression2, num % 2 is condition and for num in range(1, 11) is for_loop. In each iteration of for_loop, if condition is True, then expression1 is evaluated and added to the list, else expression2 is evaluated and added to the list.
Using List Comprehension with Nested Loops in Python
Let’s see how to use List Comprehension in case of nested loops.
In the following example, we are creating a list result which contains the elements common in two lists list1 and list2.
list1 = [1, 2, 3, 4]
list2 = [1, 20, 30, 4]
result = []
for x in list1:
for y in list2:
if x == y:
result.append(y)
print(result)
The outer for loop is iterating over the elements of list1 and the inner for loop is iterating over the elements of list2. For each iteration of the outer loop, there is a complete iteration of the inner loop, and for each iteration of the inner loop we are checking if the values of x and y are equal. If the values are equal, then it is appended to the list result.
This code can be rewritten using list comprehension as follows.
list1 = [1, 2, 3, 4]
list2 = [1, 20, 30, 4]
result = [y for x in list1 for y in list2 if x == y]
print(result)
In this example, we have two for loops. The loop for y in list2 is inside the loop for x in list1. The if statement if x == y is inside the inner for loop.
In the outer loop for x in list1, the variable x iterates over list1, and in the inner loop for y in list2, the variable y iterates over list2. In each iteration of the outer loop, the inner loop iterates over list2. In each iteration of the inner loop, if the condition x == y is True, then the variable y is added to the list.
Comparing this code to the syntax, y is expression, x == y is condition and for x in list1 and for y in list2 are two for_loops.
When to use List Comprehension
Creating lists using list comprehension is faster as compared to loops mainly because we don’t need to call the append() function in each iteration. Moreover, it takes just one line of code thus making the code more readable.
However, that doesn’t mean list comprehension should be used to replace for loop every time a list is created. If the logic is long or complex, then using list comprehension is a bad choice. This is because it can become less readable and you won’t be able add comments, thus making it difficult to debug as well.
Therefore, use list comprehension if the logic is simple and small and use for loop otherwise.
