List comprehension is an easier way of creating lists. List comprehensions are used to create lists from an iterable or a sequence. Usually, to create a list by applying some operations on each element in a sequence, we would have to use a loop. For example, the list of squares of natural numbers can be calculated as follows.
squares = []
for number in range(1, 10):
squares.append(number * number)
# Output: squares = [1, 4, 9, 16, 25, 36, 49, 64, 81]
This entire loop can be replaced with a list comprehension.
squares = [number * number for number in range(1, 10)]
This may loop like a one-liner for-loop. That is what it is. A comprehensive loop. A list comprehension has a for-loop like construct on the right side, which will generate a sequence of values and on the left side, it has the operation to perform on the values generated.
The operation on the left side is not limited to simple expressions. The following list comprehension generates a list of tuples which in the format (number, square of the number).
squares_pair = [(number, number * number) for number in range(1, 10)]
# Output: squares_pair = [(1, 1), (2, 4), (3, 9), (4, 16), (5, 25), (6, 36), (7, 49), (8, 64), (9, 81)]
Conditional Statements Inside List Comprehension
Sometimes you may want to perform certain operations on a sequence of elements only if certain conditions are met. You could use an if statement inside a for-loop to achieve this but list comprehension gives us a more comprehensive way.
Suppose you have a list of name of Game of Thrones characters and you want to filter out the Stark family members and covert their names to uppercase. You could do it in a for-loop as follows.
names = ['Arya Stark', 'Tywin Lannister', 'Sansa Stark', 'Petyr Baelish', 'Rob Stark']
starks = []
for name in names:
if name.endswith('Stark'):
starks.append(name.upper())
# Output: starks = ['ARYA STARK', 'SANSA STARK', 'ROB STARK']
You can achieve the same result with list comprehensions.
names = ['Arya Stark', 'Tywin Lannister', 'Sansa Stark', 'Petyr Baelish', 'Rob Stark']
starks = [name.upper() for name in names if name.endswith('Stark')]
# Output: starks = ['ARYA STARK', 'SANSA STARK', 'ROB STARK']
The right side which generates the sequence of elements can optionally include a conditional statement. The operation on the left side is performed only if the element satisfies the condition. It may look like a for loop and if statement combined into one line.
Nested List Comprehensions
Like nested loops, list comprehensions can also be nested. For example, if you have two sets A and B, their Cartesian Product A x B can be calculated using nested for-loops as follows.
A = [1, 2, 3]
B = ['a', 'b', 'c']
AxB = []
for a in A:
for b in B:
AxB.append((a, b))
# Output: AxB = [(1, 'a'), (1, 'b'), (1, 'c'), (2, 'a'), (2, 'b'), (2, 'c'), (3, 'a'), (3, 'b'), (3, 'c')]
Here, for illustration purposes, we are using lists to represent sets. A more correct way is to use sets and set comprehension. This is explained in the next section.
Set Comprehension
Set comprehensions are similar to list comprehensions but the result will be a set instead of a list. It is very useful when you want to remove the duplicates from the result of a comprehension. The Cartesian product example in the previous section can be converted to set comprehension as follows.
A = {1, 2, 3}
B = {'a', 'b', 'c'}
AxB = {(a, b) for a in A for b in B}
# Output: AxB = {(1, 'a'), (1, 'b'), (1, 'c'), (2, 'a'), (2, 'b'), (2, 'c'), (3, 'a'), (3, 'b'), (3, 'c')}
Dict Comprehension
Similar to lists and sets, comprehensions can be used with dictionaries too. Remember the squares example in the beginning? It would be much better if we can represent the number and its square as a key-value pair instead of a tuple.
squares = {number: number * number for number in range(1, 10)}
# Output: squares = {1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}
Instead of giving a single value, for dict comprehension we have to give a key-value pair.