Extreme Python

Logo

Extremely good python content :)

About

Articles

How to use python comprehensions

by Alan

Python comprehensions are a concise way of creating new built-in containers from existing iterables.

There are three types of containers you can create using comprehensions: list, dictionary, and set. A list comprehension1 creates a new list, a set comprehension creates a new set, and a dictionary comprehension2 creates a new dictionary.

Why the word “comprehension”?

Python uses the word “comprehension” to mean “inclusion” (borrowed from set theory). While this word confused me initially, the mathematical definition is pretty darn close to the programmatic use.

List comprehensions

Lets create a new list containing the squares of all the elements from the original list (the expression [i * i for i in numbers] is the list comprehension expression):

>>> numbers = [1, 2, 3]
>>> [i * i for i in numbers]
[1, 4, 9]

Perform the expression i * i only if i is an odd number:

>>> numbers = [1, 2, 3]
>>> [i * i for i in numbers if i % 2 == 1]
[1, 9]

Note that the variable i can be referenced throughout the entire expression.

Now lets double a value twice using nested comprehensions:

>>> [x * 2 for x in [i * 2 for i in [1, 2, 3]]]
[4, 8, 12]

Create a list of lists:

>>> numbers = [1, 2, 3]
>>> [[x * 2 for x in range(3)] for i in numbers]
[[0, 2, 4], [0, 2, 4], [0, 2, 4]]

Cartesian product (A x B) of list A and list B:

>>> A = [1, 2]
>>> B = ['A', 'B']
>>> [(i, j) for i in A for j in B]
[(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]

This is equivalent to:

>>> result = []
>>> for i in A:
...     for j in B:
...             result.append((i, j))
...
>>> result
[(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]

Takeaways

The structure of a list comprehension is pretty similar to other types of comprehensions, so you can apply most of what you know about list comprehensions to set and dictionary comprehensions as well.

Dictionary comprehensions

The syntax for dictionary comprehensions is very similar to list comprehensions.

Rather than creating new list objects from an iterable, we’re creating new dictionaries.

Create a new dictionary from a list of numbers where we map the value of the number to its square:

>>> { i:i*i for i in numbers }
{1: 1, 2: 4, 3: 9}

Just like list comprehensions, nested loops are supported:

>>> {i*j:i*i for i in numbers for j in numbers}
{1: 1, 2: 4, 3: 9, 4: 4, 6: 9, 9: 9}

Set comprehensions

Most similar syntactically to dictionary comprehensions. All you have to do is keep the braces and drop the colon :.

Create a set with all the elements in a list:

>>> numbers = [1, 1, 2, 2, 3, 3]
>>> {i for i in numbers}
{1, 2, 3}

Just for fun, lets add a predicate to exclude odd numbers:

>>> numbers = [1, 1, 2, 2, 3, 3]
>>> {i for i in numbers if i % 2 == 0}
{2}

References

  1. https://www.python.org/dev/peps/pep-0274/ 

  2. https://www.python.org/dev/peps/pep-0274/ 

tags: