List comprehensions, set comprehensions, and dictionary comprehensions in python by example

Before we learn about list comprehensions, the first step is to be familiar with lists. If you need a refresher course on lists checkout Python Lists – Learn Data Structures in Python. Let’s assume that we have a list x, and that we want to create a new list y that scales the elements of x by three. We can accomplish this by doing the following:

y = []
x = [2, 4, 6, 8, 10]
for nums in x:
    nums *= 3
    y.append(nums)

>>> y
[6, 12, 18, 24, 30]

However, you can go a step further and translate the above code into a list comprehension as follows:

y = [nums * 3 for nums in x]
>>> y
[6, 12, 18, 24, 30]

This list comprehension replicates the logic in a single line of code and is still easy to understand. Below is the pseudocode for how this list comprehension works:

y = [expression for for statement]

In English it says multiply nums by three, for every nums in the list x.

Let’s look at more code:

>>> nums = []
>>> for x in [2, 4, 6]:
...         for y in [3, 6, 9]:
...             if x != y:
...                 nums.append((x, y))
... 
>>> nums
[(2, 3), (2, 6), (2, 9), (4, 3), (4, 6), (4, 9), (6, 3), (6, 9)]

We can replicate the above logic using list comprehensions. The list comprehension should follow the same order of the for loops and if statements. As you can see in the above code there’s two loops and then a conditional statement. Below is how the translated list comprehension looks:

>>> compre = [(x,y) for x in [2,4,6] for y in [3,6,9] if x != y]
>>> compre
[(2, 3), (2, 6), (2, 9), (4, 3), (4, 6), (4, 9), (6, 3), (6, 9)]

In other words it’s like the following:

[expression for statement for statement condition]

Below is another example of a nested loop in python:

>>> for x in range(5):
...         for y in range(x):
...             print(y, end= ' ')
...     print()
... 

0 
0 1 
0 1 2 
0 1 2 3 

We can again replicate the logic with a nested list comprehension. Remember, the order of the loops in the list comprehension should be of that of the traditional loops.

>>> [y for x in range(5) for y in range(x)]
[0, 0, 1, 0, 1, 2, 0, 1, 2, 3]

Let’s translate another example into a list comprehension. Let’s suppose that you wanted to get all of the even numbers of a list named nums:

nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

We could use the traditional for loop as follows:

>>> for x in nums:
...         if x % 2 == 0:
...             print(x, end=' ')
... 
2 4 6 8 10 

However, we can translate the above statement into a list comprehension as follows:

>>> [even  for even  in nums if even % 2 == 0
[2, 4, 6, 8, 10]

As you can see, the list comprehension follows the order of the for and if statements of the original code snippet. Next, let’s assume that we wanted to create a list of length 10 filled with random numbers in the range of 1…100. One way we could do that is as follows:

>>> import random
>>> rnums = []
>>> for x in range(10):
...         rnums.append(random.randint(1, 100))
... 
>>> rnums
[27, 49, 70, 25, 74, 28, 61, 65, 9, 12]

But wait! With our new friend the list comprehension, we can replicate the logic in a more concise manner as shown below:

>>> [random.randint(1, 100) for x in range(10)]
[57, 90, 45, 24, 10, 18, 59, 94, 75, 89]

Set comprehensions

Just like with list comprehensions you can do set comprehensions. To get a refresher course on sets in python read Learn How to Use Sets in Python. The following syntax creates a set in python:

>>> nums = set()
>>> nums.add(2)
>>> nums.add(4)
>>> nums.add(6)
>>> nums.add(8)
>>> nums.add(10)
>>> nums
{2, 4, 6, 8, 10}

The above can be replicated using a set comprehension:

>>> even = {x for x in range(1, 10) if x % 2 == 0}
>>> sorted(even)
[2, 4, 6, 8]

The syntax of a set comprehensions closely resembles that of a list comprehension except that curly braces {} are used. Let’s look at the following code which tests for uniqueness:

>>> import random
>>> rnums = []
>>> for x in range(15):
...         rnums.append(random.randint(1, 10))
... 

The above code simply creates an empty list named rnums of length 15, and then randomly fills it with random integers in the range of 1-10.

>>> rnums
[7, 7, 6, 10, 5, 8, 2, 7, 6, 1, 5, 5, 3, 3, 4]
>>> unique = set()
>>> for x in rnums:
...         unique.add(x)
... 
>>> unique
{1, 2, 3, 4, 5, 6, 7, 8, 10}

The above code iterates over the random list and then add it’s elements to a set. Since sets contains unique entries the set will only return non-duplicated entries. We can reduce the above code by using a set comprehension instead of the for loop. Here’s the above code modified to use a set comprehension:

>>> unique = {x for x in rnums}
>>> sorted(unique)
[1, 2, 3, 4, 5, 6, 7, 8, 10]

The sorted() function is called because sets are unordered.

Dictionary comprehensions

Before learning about dictionary comprehensions you should have an idea about dictionaries. Here’s a blog post I made that reveals the inside scoop about dictionaries in python.

Since dictionaries and sets both use curly braces ({}) a dictionary comprehension needs a way to distinguish itself. The following shows the general syntax of a dictionary comprehension in python:

{key: value for (key, value) in iterable}

Below is one example of a dictionary comprehension in python:

>>> {a: 0 for a in range(5)}
{0: 0, 1: 0, 2: 0, 3: 0, 4: 0}

You can use the enumerate() function to count the number of elements.

>>> {i:j for i, j in enumerate(range(5))}
{0: 0, 1: 1, 2: 2, 3: 3, 4: 4}
>>> letters = ['a', 'b', 'c', 'd']
>>> {letters: num for num, letters  in enumerate(letters)}
{'a': 0, 'b': 1, 'c': 2, 'd': 3}

Here’s another dictionary comprehension example:

>>> letters = ['a', 'b', 'c', 'd']
>>> {l:l.capitalize() for l in letters}
{'a': 'A', 'b': 'B', 'c': 'C', 'd': 'D'}
============================================================================ Want to learn how to use Python's most popular IDE Pycharm? In the free pdf guide "Getting the Hang of PyCharm" you'll learn all of the amazing features in PyCharm along with how to get started with data science. Subscribe to the Purcell Consult newsletter and get started A.S.A.P.