I’m a fan of The Zen of Python, which says
There should be one — and preferably only one — obvious way to do it.
But in Python, there are in fact many ways to achieve the same objective. Of course, some ways are more elegant than others and in most cases, it should be obvious which way is better.
We are going to look at list comprehensions, and how they can replace for loops, map()
and filter()
to create powerful functionality within a single line of Python code.
Basic List Comprehension
Say I want to create a list of numbers from 1 to 10. I could do
numbers = []
for i in range(1, 11):
numbers.append(i)
and I would get
>>> numbers
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
But using list comprehension, this could be accomplished within a single line
>>> numbers = [i for i in range(1, 11)]
>>> numbers
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
This is the basic syntax of list comprehension: [expression for element in iterable]
. Here, the iterable is range(1, 11)
, the element is i
and the expression is i
. This is equivalent to the for loop we used earlier: we add i
to the list where i
is a number from 1 to 11.
map()
The map()
function is often used to apply a function on each element in an iterable. Pass in a function and an iterable, and map()
will create an object containing the results of passing each element into the function.
For example, say I want to create a list of squares from the numbers
list we created earlier. We could do
squares = []
for num in numbers:
squares.append(num ** 2)
and we will get
>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
Or we could use map()
like so
>>> squares = list(map(lambda x: x ** 2, numbers))
>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
Here, we pass each of the elements in numbers
into the lambda function (which is just an easy way to create functions if you are lazy). The output of putting each number x
into the function lambda x: x ** 2
will be the square of the number. Using list()
we turn the map object into a list.
Replacing map() With List Comprehension
Using list comprehension, we could simply do
>>> squares = [num ** 2 for num in numbers]
>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
This will pass each number num
into the expression num ** 2
and create a new list where the elements are simply the squares of each number in numbers
.
filter()
The filter()
function is used to create a subset of an existing list, based on a condition. Pass in a function and an iterable, and filter()
will create an object containing all elements where the function evaluates to True
.
For example, say I want to get a list of all odd numbers and a list of all even numbers from the numbers
list we created earlier. We could do
odd_numbers = []
even_numbers = []
for num in numbers:
if num % 2 == 1:
odd_numbers.append(num)
elif num % 2 == 0:
even_numbers.append(num)
and we would get
>>> odd_numbers
[1, 3, 5, 7, 9]
>>> even_numbers
[2, 4, 6, 8, 10]
Or we could use filter()
like so
odd_numbers = list(filter(lambda x: x % 2 == 1, numbers))
even_numbers = list(filter(lambda x: x % 2 == 0, numbers))
Here, all the numbers will be passed into the lambda function and if x % 2 == 1
is True
, the number will be included in odd_numbers
. Likewise, if x % 2 == 0
is True
, the number will be included in even_numbers
.
Replacing filter() With List Comprehension
Using list comprehension, we could simply do
odd_numbers = [num for num in numbers if num % 2 == 1]
even_numbers = [num for num in numbers if num % 2 == 0]
Here, we are using a conditional. The syntax for this is [expression for element in iterable (if conditional)]
. This is equivalent to the for loop we used earlier — if the condition num % 2 == 1
is met, num
will be added to odd_numbers
, and num
is an element in the numbers
iterable.
Nested Comprehensions
Say we want to create a matrix. This would involve creating nested lists. Using a for loop, we can do the following
matrix = []
for i in range(5):
row = []
for j in range(5):
row.append(i * 5 + j)
matrix.append(row)
and we would get
>>> matrix
[
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]
]
Using a nested comprehension, we could simply do
matrix = [[i * 5 + j for j in range(5)] for i in range(5)]
The outer list comprehension [... for i in range(5)]
creates 5 rows, while the inner list comprehension [... for j in range(5)]
creates 5 columns.
Dictionary Comprehensions
You could use dictionary comprehensions too. For example, if I want to create a dictionary that maps each number in numbers
to their corresponding square, we could do
num_to_square = {}
for num in numbers:
num_to_square[num] = num ** 2
and we would get
>>> num_to_square
{
1: 1,
2: 4,
3: 9,
4: 16,
5: 25,
6: 36,
7: 49,
8: 64,
9: 81,
10: 100
}
Or we could use a dictionary comprehension
>>> num_to_square = {num: num ** 2 for num in numbers}
>>> num_to_square
{
1: 1,
2: 4,
3: 9,
4: 16,
5: 25,
6: 36,
7: 49,
8: 64,
9: 81,
10: 100
}
You could even use a dictionary comprehension within a list comprehension or vice versa.
>>> num_to_square = [{num: num ** 2 for num in numbers} for numbers in
[odd_numbers, even_numbers]]
>>> num_to_square
[
{1: 1, 3: 9, 5: 25, 7: 49, 9: 81},
{2: 4, 4: 16, 6: 36, 8: 64, 10: 100}
]
Why Use List Comprehension?
- Cleaner, clearer code
- Slightly faster than
map()
andfilter()
- Generally considered more ‘pythonic’
But hey, at the end of the day, the choice is yours. List comprehension is just another way to do the same things, and whether something is ‘cleaner’ and ‘clearer’ is largely subjective. However, most people would agree that list comprehension is a more standard way to do things.
Conclusion
That’s it! Now you know how to use list comprehensions in Python.
If you have any questions, feel free to comment down below.
Thanks for reading, and follow me for more stories like this in the future.
Top comments (0)