List Comprehension in Python

Today we’ll be talking about the list comprehension in Python.

A list comprehension is used to create a list. It can be used instead of a couple of functions we discussed before: the lambda function, the map function, and the filter function. They cover the functionality of them all.

If you want to learn more about the lambda function, the map function and the filter function, I have an article on each of them, so feel free to read them.

In a list comprehension we tell the program to do something to each element of an iterable, like a string, list, range or any other iterable.

I also have an article on iterables and iterators, so you can read it, too, if you want.

The Three Approaches

Let’s start with something simple. Suppose we need a list of even numbers up to 20. We can accomplish this with a for loop, the filter function with a lambda or a list comprehension. Let’s compare the three solutions. Here’s the version with a for loop:

>>> even_numbers = []
>>> for number in range(21):
...     if number % 2 == 0:
...         even_numbers.append(number)
... 
>>> even_numbers
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

Now, the version with the filter function with a lambda:

>>> even_numbers = list(filter(lambda number: number % 2 == 0, range(21)))
>>> even_numbers
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

Finally, the most concise and readable version with a list comprehension:

>>> even_numbers = [number for number in range(21) if number % 2 == 0]
>>> even_numbers
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]

Although we get the same results in the three cases, the for loop is slower than the filter function or the list comprehension, which may matter in more complex examples.

The Loop Variable in List Comprehensions

In list comprehensions the loop variable is local and has no effect on global variables with the same name:

>>> x = "Hello"  # global x

>>> odd_numbers = [x for x in range(20) if x % 2]  # local x
>>> print(odd_numbers)
[1, 3, 5, 7, 9, 11, 13, 15, 17, 19]

>>> print(x)  # global x
Hello

An Example with Strings

And here’s an example in which the program adds “.pdf” at the end of each file name in the files list:

files = ["My cat", "Sister smoking in cellar", "Disappearing dinner"]
full_names = [file + (".pdf") for file in files]
print(full_names)

Here’s the output:

['My cat.pdf', 'Sister smoking in cellar.pdf', 'Disappearing dinner.pdf']

An Example with Numbers

And now an example with numbers. First we turn the numbers into binary numbers, and in the other comprehension we use a formula to change the values.

decimal_numbers = [4, 7, 8, 11]
binary_numbers = [bin(n) for n in decimal_numbers]
print(binary_numbers)
encoded = [2 * (n - 1) for n in decimal_numbers]
print(encoded)

Here’s the output:

['0b100', '0b111', '0b1000', '0b1011']
[6, 12, 14, 20]

Multiple Loop Variables

We can also use multiple loop variables in a list comprehension. Here’s an example that will print all the Pythagorean triples such that each integer will be less than 100. What’s a Pythagorean triple? It’s a group of three positive integers (a, b, c) such that a² + b² = c². Here’s the code:

pythagorean_triples = [(x, y, z) for x in range(1, 100) for y in range(x, 100) for z in range(y, 100) if x ** 2 + y ** 2 == z ** 2]

which is more readable when rewritten like so:

pythagorean_triples = [(x, y, z) 
                       for x in range(1, 100) 
                       for y in range(x, 100) 
                       for z in range(y, 100) 
                       if x ** 2 + y ** 2 == z ** 2]

And now let’s print the triples:

print(pythagorean_triples)

And here’s the output. The list comprehension returns a list of 3-tuples:

[(3, 4, 5), (5, 12, 13), (6, 8, 10), (7, 24, 25), (8, 15, 17), (9, 12, 15), (9, 40, 41), (10, 24, 26), (11, 60, 61), (12, 16, 20), (12, 35, 37), (13, 84, 85), (14, 48, 50), (15, 20, 25), (15, 36, 39), (16, 30, 34), (16, 63, 65), (18, 24, 30), (18, 80, 82), (20, 21, 29), (20, 48, 52), (21, 28, 35), (21, 72, 75), (24, 32, 40), (24, 45, 51), (24, 70, 74), (25, 60, 65), (27, 36, 45), (28, 45, 53), (30, 40, 50), (30, 72, 78), (32, 60, 68), (33, 44, 55), (33, 56, 65), (35, 84, 91), (36, 48, 60), (36, 77, 85), (39, 52, 65), (39, 80, 89), (40, 42, 58), (40, 75, 85), (42, 56, 70), (45, 60, 75), (48, 55, 73), (48, 64, 80), (51, 68, 85), (54, 72, 90), (57, 76, 95), (60, 63, 87), (65, 72, 97)]

For better readability let’s format the output:

for triple in pythagorean_triples:
    print(f"a = {triple[0]}, b = {triple[1]}, c = {triple[2]}")

And here’s the output now:

a = 3, b = 4, c = 5
a = 5, b = 12, c = 13
a = 6, b = 8, c = 10
a = 7, b = 24, c = 25
a = 8, b = 15, c = 17
a = 9, b = 12, c = 15
a = 9, b = 40, c = 41
a = 10, b = 24, c = 26
a = 11, b = 60, c = 61
a = 12, b = 16, c = 20
a = 12, b = 35, c = 37
a = 13, b = 84, c = 85
a = 14, b = 48, c = 50
a = 15, b = 20, c = 25
a = 15, b = 36, c = 39
a = 16, b = 30, c = 34
a = 16, b = 63, c = 65
a = 18, b = 24, c = 30
a = 18, b = 80, c = 82
a = 20, b = 21, c = 29
a = 20, b = 48, c = 52
a = 21, b = 28, c = 35
a = 21, b = 72, c = 75
a = 24, b = 32, c = 40
a = 24, b = 45, c = 51
a = 24, b = 70, c = 74
a = 25, b = 60, c = 65
a = 27, b = 36, c = 45
a = 28, b = 45, c = 53
a = 30, b = 40, c = 50
a = 30, b = 72, c = 78
a = 32, b = 60, c = 68
a = 33, b = 44, c = 55
a = 33, b = 56, c = 65
a = 35, b = 84, c = 91
a = 36, b = 48, c = 60
a = 36, b = 77, c = 85
a = 39, b = 52, c = 65
a = 39, b = 80, c = 89
a = 40, b = 42, c = 58
a = 40, b = 75, c = 85
a = 42, b = 56, c = 70
a = 45, b = 60, c = 75
a = 48, b = 55, c = 73
a = 48, b = 64, c = 80
a = 51, b = 68, c = 85
a = 54, b = 72, c = 90
a = 57, b = 76, c = 95
a = 60, b = 63, c = 87
a = 65, b = 72, c = 97

Cross Products

We can also use a list comprehension for the cross product of two sets. The cross product, also known as Cartesian product A x B, is a set of all the pairs of elements of which one element belongs to set A and the other to set B. So, what we get are all the possible combinations. Here’s an example. First let’s define the two sets (actually here these are lists) that we want to work on. We want to combine simple verbs with particles to make phrasal verbs. Phrasal verbs are verbs usually consisting of two words, like ‘make up’, ‘show off’, “turn on” and many, many others.

>>> verbs = ["get", "set", "give"]
>>> particles = ["up", "in", "on", "off"]
>>> phrasal_verbs = [(verb, particle) for verb in verbs for particle in particles]

Now, if we print the list created by means of the list comprehension, we’ll see that it’s a list of 2-tuples of comma-separated parts of phrasal verbs:

>>> print(phrasal_verbs)
[('get', 'up'), ('get', 'in'), ('get', 'on'), ('get', 'off'), ('set', 'up'), ('set', 'in'), ('set', 'on'), ('set', 'off'), ('give', 'up'), ('give', 'in'), ('give', 'on'), ('give', 'off')]

Now we can use the phrasal_verbs list in another list comprehension to get a list of strings, each being a full phrasal verb:

>>> phrasals = [f"{verb[0]} {verb[1]}" for verb in phrasal_verbs]

Let’s check it out :

>>> print(phrasals)
['get up', 'get in', 'get on', 'get off', 'set up', 'set in', 'set on', 'set off', 'give up', 'give in', 'give on', 'give off'] 

Works.

Although we should keep list comprehensions simple, we could achieve the same in just one step:

>>> verbs = ["get", "set", "give"]
>>> particles = ["up", "in", "on", "off"]
>>> phrasals = [f"{verb[0]} {verb[1]}" for verb in 
...             [(verb, particle) for verb in verbs for particle in particles]]
... 
>>> phrasals
['get up', 'get in', 'get on', 'get off', 'set up', 'set in', 'set on', 'set off', 'give up', 'give in', 'give on', 'give off']

Such nested comprehensions let you save some typing, but they also obfuscate your code, so it’s usually better to keep them simple.

Examples of List Comprehensions in Interactive Mode

Here are some examples of how we can quickly use comprehensions in interactive mode:

>>> numbers = [-5, -3, 0, 4, 6]

– squared numbers:

>>> [n ** 2 for n in numbers]
[25, 9, 0, 16, 36]

– odd numbers:

>>> [n for n in numbers if n % 2]
[-5, -3]

– absolute values of numbers:

>>> [abs(n) for n in numbers]
[5, 3, 0, 4, 6]

– strings containing number characters:

>>> ["number " + str(n) for n in numbers]
['number -5', 'number -3', 'number 0', 'number 4', 'number 6']

– numbers with their opposites:

>>> [(n, -n) for n in numbers]
[(-5, 5), (-3, 3), (0, 0), (4, -4), (6, -6)]

– Boolean representations of numbers:

>>> [bool(n) for n in numbers]
[True, True, False, True, True]

Here’s the video version of the article: