List comprehension is an easier way of creating lists. Typically to create a list by applying some operations on each element of 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 by 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 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 values. 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)]

Conditionals 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 could achieve the same result with list comprehensions but in a more concise way.

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 illustrations purposes, we are using lists to represent sets. A more correct way is to use set and set comprehension which is explained in the next section.

Set Comprehension

Set comprehensions are similar to list comprehensions but result will be a set instead of a list. It is very useful when you want 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 also. 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.


Posted on
Category: Python
Tags: Python, List Comprehensions, Set Comprehensions, Dict Comprehensions