Comprehensions is an intuitive way to create lists, dicts, and sets. It is also a good alternative for using loops, maps, or filters.
In this article, we will learn how to use comprehension, and you will see how it will make the codes more elegant and readable.
List Comprehension
First, let's see how we can create a list using comprehension and a for-loop.
Comprehension with a for-loop
Scenario #1: We want n for each n in nums
Create a list with a for-loop:
nums = [1,2,3,4,5,6,7,8,9]
result_list = []
for n in nums:
result_list.append(n)
print(result_list)
# output ->
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
Accomplish the same result with list comprehension:
nums = [1,2,3,4,5,6,7,8,9]
result_list = [num for num in nums]
print(result_list)
# output ->
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
As you can see, the codes are much simpler to write with comprehension than using a normal for-loop. It is also better to understand. When you read the code within the brackets, you feel like you are reading the scenario above.
# We want 'n' for each 'n' in 'nums'
result_list = [num for num in nums]
We can also do some calculations, such as squaring a list of the number:
nums = [1,2,3,4,5,6,7,8,9]
result_list = [num**2 for num in nums]
print(result_list)
# output ->
# [1, 4, 9, 16, 25, 36, 49, 64, 81]
If you know how to use map(), then you can do something similar with it:
nums = [1,2,3,4,5,6,7,8,9]
result_list = map(lambda n: n**2, nums)
print(result_list)
# output ->
# [1, 4, 9, 16, 25, 36, 49, 64, 81]
It is also a one-liner, but the list comprehension is more readable and easier to understand. It demonstrates that we can use list comprehension as an alternative to map().
Comprehension with if statement
Based on the previous example, let's add a logical statement.
Scenario #2: We want n for each n if n is even.
Using a normal for-loop style:
nums = [1,2,3,4,5,6,7,8,9]
result_list = []
for n in nums:
# Add an IF statement here
if n % 2 == 0:
result_list.append(n)
print(result_list)
# output ->
# [2, 4, 6, 8]
For list comprehension, we can add the IF statement at the tail of the expression within the brackets:
nums = [1,2,3,4,5,6,7,8,9]
result_list = [num for num in nums if num % 2 == 0]
print(result_list)
# output ->
# [2, 4, 6, 8]
Just like how we used the map function, we can use the filter function here to do achieve the same result:
nums = [1,2,3,4,5,6,7,8,9]
result_list = filter(lambda n: n % 2 == 0, nums)
print(result_list)
# output ->
# [2, 4, 6, 8]
What the filter function above does is to keep the even number of the list and ignore others that are not.
Okay, now we understand the basic. Let's move on to a more complex example.
Comprehension with two (or mutiple) for-loops
Scenario #3: We want a (letter, num) tuple for each letter in 'ABCD' and each numbe in '1234'
Normal for-loop style:
result_list = []
for letter in 'ABCD':
for num in '1234':
result_list.append((letter, num))
print(result_list)
# output ->
# [('A', '1'), ('A', '2'), ('A', '3'), ('A', '4'), ('B', '1'), ('B', '2'), ('B', '3'), ('B', '4'), ('C', '1'), ('C', '2'), ('C', '3'), ('C', '4'), ('D', '1'), ('D', '2'), ('D', '3'), ('D', '4')]
For list comprehension, we can write the for-loop one after another within the brackets:
result_list = [(letter, num) for letter in 'ABCD' for num in '1234']
print(result_list)
# output ->
# [('A', '1'), ('A', '2'), ('A', '3'), ('A', '4'), ('B', '1'), ('B', '2'), ('B', '3'), ('B', '4'), ('C', '1'), ('C', '2'), ('C', '3'), ('C', '4'), ('D', '1'), ('D', '2'), ('D', '3'), ('D', '4')]
Dictionary Comprehension
We can also use comprehension to create a dictionary.
Scenario #4: Create a dictionary {'id': 'fruit'} for each (id, fruit) in zip(ids, fruits)
ids = [1,2,3,4,5]
fruits = ['apple', 'orange', 'mango', 'banana', 'strawberry']
Let's do it with a for-loop:
dict = {}
for id, fruit in zip(ids, fruits):
dict[id] = fruit
print(dict)
# output ->
# {1: 'apple', 2: 'orange', 3: 'mango', 4: 'banana', 5: 'strawberry'}
This is how we do it with dictionary comprehension:
dict = {id: fruit for id, fruit in zip(ids, fruits)}
print(dict)
# output ->
# {1: 'apple', 2: 'orange', 3: 'mango', 4: 'banana', 5: 'strawberry'}
Like the list comprehension, we can add an IF statement to the dictionary comprehension:
dict = {id: fruit for id, fruit in zip(ids, fruits) if id != 2}
print(dict)
# output ->
# {1: 'apple', 3: 'mango', 4: 'banana', 5: 'strawberry'}
Set Comprehension
The comprehension works for creating a set too!
Let's create a set with a for-loop:
nums = [2,3,2,4,25,6,4,2,5,6]
my_set = set()
for n in nums:
my_set.add(n)
print(my_set)
# output -> {2, 3, 4, 5, 6, 25}
To use comprehension to create a set, we write it like the list comprehension. Instead of using brackets [...], we use braces {...} :
nums = [2,3,2,4,25,6,4,2,5,6]
my_set = {num for num in nums}
print(my_set)
# output -> {2, 3, 4, 5, 6, 25}
Generator Comprehension
Finally, let's demonstrate how to create a generator with comprehension.
nums = [1,2,3,4,5,6,7,8,9]
def create_gen(nums):
for n in nums:
yield n ** 2
my_gen = create_gen(nums)
for n in my_gen:
print(n)
# output:
# 1
# 4
# 9
# ...
For generator comprehension, we use parenthesis (...) to wrap the statement:
nums = [1,2,3,4,5,6,7,8,9]
my_gen = (num ** 2 for num in nums)
for n in my_gen:
print(n)
# output:
# 1
# 4
# 9
# ...That's it! By using comprehensions properly, we not only can make our codes easier to read but also make them easier to write.