List comprehensions provide a way of writing for loops more concisely. They can be useful when you want to create new lists from existing lists or iterables. For example, you can use list comprehensions to create a list of squares from a list of numbers. Similarly, you could also use some conditions on a list of numbers so that the new list you create is a subset of the original list.
- Loops and List Comprehensions
- Conditions in List Comprehensions
- Efficient List Comprehensions With the Walrus Operator
- List Comprehensions for Nested Loops
- Nested List Comprehensions
- Set and Dictionary Comprehensions
Keep in mind that you cannot write every for loop as a list comprehension. One more thing: the name "list comprehensions" can be a bit confusing because it seems to suggest that the comprehensions are only meant for working with lists. In reality, the word "list" in list comprehensions is used to indicate that you can loop over any iterable in Python, and the end product will be a list.
Loops and List Comprehensions
Basic list comprehensions that don't use any conditions have the following form:
1 |
[<the_expression> for <the_element> in <the_iterable>] |
Let's begin by writing a very basic for loop to list the first 15 multiples of 5. First, you need to create an empty list. Then, you have to iterate over a range of numbers and multiply them by 5. The new sequence of numbers that you get will consist of multiples of 5.
1 |
multiples = [] |
2 |
|
3 |
for n in range(1,16): |
4 |
multiples.append(n*5) |
5 |
|
6 |
print(multiples) |
7 |
# [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75]
|
The above for loop basically has the following structure:
1 |
for <the_element> in <the_iterable>: |
2 |
<the_expression> |
If you compare it with the list comprehension form that you read earlier, you can see that <the_element> is n, <the_iterable> is range(1,16), and <the_expression> is n*5. Putting these values in the list comprehension will give us the following result:
1 |
multiples = [n*5 for n in range(1,15)] |
2 |
|
3 |
multiples
|
4 |
# [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70]
|
Similarly, you can also get a list with the cube of given numbers like this:
1 |
cubes = [n**3 for n in range(1,16)] |
2 |
|
3 |
print(cubes) |
4 |
#[1, 8, 27, 64, 125, 216, 343, 512, 729, 1000, 1331, 1728, 2197, 2744, 3375]
|
Conditions in List Comprehensions
You can also use an if condition to filter out certain values from the final list. In this case, the list comprehension takes the following form:
1 |
[<the_expression> for <the_element> in <the_iterable> if <the_condition>] |
A basic example of this type of comprehension would be to get all the even numbers in a given range. A for loop to do this task will look like this:
1 |
evens = [] |
2 |
|
3 |
for n in range(1,21): |
4 |
if n%2 == 0: |
5 |
evens.append(n) |
6 |
|
7 |
print(evens) |
8 |
# [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
|
The same thing could also be accomplished by using the following list comprehension:
1 |
evens = [n for n in range(1,21) if n%2 == 0] |
2 |
|
3 |
print(evens) |
4 |
# [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
|
A more complex example of using list comprehensions would be adding .. if .. else .. conditional expressions inside them.
In this case, the order in which you lay out the statements inside the list comprehension will be different from the usual if conditions. When you only have an if condition, the condition goes to the end of the comprehension. However, in the case of an .. if .. else .. expression, the positions of the for loop and the conditional expression are interchanged. The new order is:
1 |
[<the_expression> if <the_condition> else <other_expression> for <the_element> in <the_iterable>] |
Let's begin by writing the verbose .. if .. else .. conditional expression to get squares of even numbers and cubes of odd numbers in a given range.
1 |
squares_cubes = [] |
2 |
|
3 |
for n in range(1,16): |
4 |
if n%2 == 0: |
5 |
squares_cubes.append(n**2) |
6 |
else: |
7 |
squares_cubes.append(n**3) |
8 |
|
9 |
print(squares_cubes) |
10 |
# [1, 4, 27, 16, 125, 36, 343, 64, 729, 100, 1331, 144, 2197, 196, 3375]
|
The above conditional expression has the following structure:
1 |
for <the_element> in <the_iterable>: |
2 |
if <the_condition>: |
3 |
<the_expression> |
4 |
else: |
5 |
<other_expression> |
Putting the corresponding values in the right places will give you the following list comprehension:
1 |
squares_cubes = [n**2 if n%2 == 0 else n**3 for n in range(1,16)] |
2 |
|
3 |
print(squares_cubes) |
4 |
# [1, 4, 27, 16, 125, 36, 343, 64, 729, 100, 1331, 144, 2197, 196, 3375]
|
Efficient List Comprehensions With the Walrus Operator
In the previous section, we created a list with a mix of squares and cubes of natural numbers. What if the criteria for storing the squares or cubes were not the odd or even values of the numbers but the divisibility of their square by 4 or 9? The following list comprehension will give back all numbers that meet our desired criteria.
1 |
square_cubes = [n**2 if n**2 % 9 == 0 or n**2 % 4 == 0 else n**3 for n in range(1, 11) ] |
2 |
|
3 |
# [1, 4, 9, 16, 125, 36, 343, 64, 81, 100]
|
4 |
print(square_cubes) |
One thing that you might have noticed is that we calculated the square three times in the above example. What if we could store the result somewhere so that the calculation has to be done only once? This is possible with the help of the walrus operator in Python, which was introduced in version 3.8.
The walrus operator is denoted by :=, and it allows you to calculate the value of an expression and assign that value to a variable at the same time. We can rewrite an optimized version of our code using the walrus operator as follows:
1 |
square_cubes = [res if (res := n**2) % 9 == 0 or res % 4 == 0 else n**3 for n in range(1, 11)] |
2 |
|
3 |
# [1, 4, 9, 16, 125, 36, 343, 64, 81, 100]
|
4 |
print(square_cubes) |
Here, we stored the result of n**2 in the res variable and then reused the variable in all the places where we were previously calculating n**2.
List Comprehensions for Nested Loops
You can also use nested loops within a list comprehension. There is no limit on the number of for loops that you can put inside a list comprehension. However, you have to keep in mind that the order of the loops should be the same in both the original code and the list comprehension. You can also add an optional if condition after each for loop. A list comprehension with nested for loops will have the following structure:
1 |
[ <the_expression> for <element_a> in <iterable_a> (optional if <condition_a>) |
2 |
for <element_b> in <iterable_b> (optional if <condition_b>) |
3 |
for <element_c> in <iterable_c> (optional if <condition_c>) |
4 |
... and so on ...] |
The following examples should make everything clearer. There are two nested loops, and multiplying them together gives us multiplication tables.
1 |
multiplications = [] |
2 |
|
3 |
for i in range(1, 4): |
4 |
for n in range(1, 11): |
5 |
multiplications.append(i*n) |
6 |
|
7 |
print(multiplications) |
8 |
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
|
These nested for loops can be rewritten as:
1 |
for <element_a> in <iterable_a>: |
2 |
for <element_b> in <iterable_b>: |
3 |
<the_expression> |
Once you have written the loop in this form, converting it to a list comprehension is easy:
1 |
multiplications = [i*n for i in range(1,4) for n in range(1,11)] |
2 |
|
3 |
print(multiplications) |
4 |
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
|
You can also use a similarly written list comprehension to flatten a list of lists. The outer for loop iterates through individual lists and stores them in the row variable. The inner for loop will then iterate through all the elements in the current row. During the first iteration, the row variable has the value [1, 2, 3, 4]. The second loop iterates through this list or row and appends all those values to the final list.
1 |
matrix = [ |
2 |
[1, 2, 3, 4], |
3 |
[5, 6, 7, 8], |
4 |
[9, 10, 11, 12], |
5 |
]
|
6 |
|
7 |
flatten = [n for row in matrix for n in row] |
8 |
|
9 |
print(flatten) |
10 |
#[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
|
Nested List Comprehensions
Nested list comprehensions may sound similar to list comprehensions with nested loops, but they are very different. In the first case, you were dealing with loops within loops. In this case, you will be dealing with list comprehensions within list comprehensions. A good example of using nested list comprehensions would be transposing the matrix from the previous section.
Without a list comprehension expression, you would need to use two for loops to transpose the matrix.
1 |
matrix = [ |
2 |
[1, 2, 3, 4], |
3 |
[5, 6, 7, 8], |
4 |
[9, 10, 11, 12], |
5 |
]
|
6 |
|
7 |
transpose = [] |
8 |
|
9 |
for i in range(4): |
10 |
temp = [] |
11 |
for row in matrix: |
12 |
temp.append(row[i]) |
13 |
transpose.append(temp) |
14 |
|
15 |
print(transpose) |
16 |
# [[1, 5, 9], [2, 6, 10], [3, 7, 11], [4, 8, 12]]
|
The outer loop iterates through the matrix four times because there are four columns in it. The inner loop iterates through the elements in the current row one index at a time and appends it to a temporary list called temp. The temp list is then appended as a row to the transposed matrix. In the case of nested list comprehensions, the outermost loop comes at the end, and the innermost loop comes at the beginning.
Here is the above code in the form of a list comprehension:
1 |
matrix = [ |
2 |
[1, 2, 3, 4], |
3 |
[5, 6, 7, 8], |
4 |
[9, 10, 11, 12], |
5 |
]
|
6 |
|
7 |
transpose = [[row[n] for row in matrix] for n in range(4)] |
8 |
|
9 |
print(transpose) |
10 |
# [[1, 5, 9], [2, 6, 10], [3, 7, 11], [4, 8, 12]]
|
Another way to look at this is by considering the structure of list comprehensions that replace the basic for loops that you learned about at the beginning of the article.
1 |
[<the_expression> for <the_element> in <the_iterable>] |
If you compare it with the nested list comprehension above, you will see that <the_expression> in this case is actually another list comprehension: [row[n] for row in matrix]. This nested list comprehension itself is in the form of a basic for loop.
Set and Dictionary Comprehensions
The concept of list comprehensions is applicable to set and dictionary comprehensions as well. As you know, dictionaries in Python are used to store key-value pairs. Earlier, we used list comprehensions to store the squares of different numbers. Let's modify the code a little bit to store numbers and their squares or cubes as key-value pairs.
1 |
squares_cubes = {n: n**2 if n%2 == 0 else n**3 for n in range(1,11)} |
2 |
|
3 |
print(squares_cubes) |
4 |
# {1: 1, 2: 4, 3: 27, 4: 16, 5: 125, 6: 36, 7: 343, 8: 64, 9: 729, 10: 100}
|
We can also create sets using similar code for set comprehensions. However, you should remember that sets are useful for creating unordered collections. So they might not make much sense for storing squares or cubes if you want to access them sequentially.
In the code below, we used set comprehension to create a set of random numbers that are not multiples of 9.
1 |
import random |
2 |
|
3 |
non_multiples = {n for n in random.sample(range(0, 1001), 20) if n not in range(0, 1001, 9)} |
4 |
|
5 |
# {3, 165, 807, 574, 745, 266, 616, 44, 12, 910, 336, 145, 755, 179, 25, 796, 926}
|
6 |
print(non_multiples) |
Final Thoughts
I hope this tutorial helped you understand what list comprehensions are and how to use them in place of basic for loops to write concise and slightly faster code while creating lists.
Another thing that you should keep in mind is the readability of your code. Creating list comprehensions for nested loops will probably make the code less readable. In such cases, you can break down the list comprehension into multiple lines to improve readability.
Additionally, don’t hesitate to see what we have available for sale and for study on Envato Market, and don't hesitate to ask any questions and provide your valuable feedback using the forum.
