Python List Comprehension + Set + Dict Comprehension

Kevin Jacobs

Nov 16, 2017 • 6 min read

Python List Comprehension support is great for creating readable but compact code for representing mathematical ideas. In this blog post, the concept of list, set and dictionary comprehensions are explained and a few examples in Python are given.

Maybe slightly different but also interesting is this blog post which explains Matplotlib for Python step-by-step. Something completely different: if you are looking for a Pythonic job, but don’t know how to answer Python interview questions, take a look at this article.

Introduction

In mathematics, in the field of logic and set theory, there is a natural way of constructing lists. For example, one can construct the set of all primes by defining the following:

in square bracket notation. Or it is possible to construct less concise sets:

. The reader can guess that it is about square numbers, but it is not immediately clear to the reader since the set is not fully specified. It is also possible to create an empty list:

. All of these concepts are found in the Python language. The next few sections explain Python list comprehension, Python set comprehension and Python dictionary comprehension.

Basic concepts: List versus Set versus Dict

In Python, there are three main concepts for element containers. The first concept is a list. A list is a ordered collection of items which may contain duplicate items. An example is the following:

["apple", "apple", "orange"]

This list clearly contains two apples and two oranges. A set is an unordered collection of unique items. This also closely corresponds to the mathematical way of creating a list. If we convert the previous list to a set, it would result into the following set:

{"apple", "orange"}

Notice that since it is unordered, it is exactly the same as the following set:

{"orange", "apple"}

It is not possible to construct a set with two or more oranges or two or more apples since every item is unique in a set.

A dict is a special item container. It is a key-value container (list of tuples) in which the keys are unique. One example of a dict is the following:

{"apples": 2, "bananas": 10, "oranges": True}

Python List Comprehension

Suppose we have two lists, one containing even numbers and one containing numbers divisible by 3:

We can create another list which contains numbers both even and divisible by 3 easily:

C = \{c | c \mbox{ is even and divisible by 3}\} = \{c | c \in A \wedge c \in B\}

In Python code, this is the same as the following:

# Create the lists
A = [a for a in range(20) if a % 2 == 0]
B = [b for b in range(20) if b % 3 == 0]
C = [c for c in range(20) if c in A and c in B]

# Print out the result
print('A: ', A)
print('B: ', B)
print('C: ', C)

This results into the following:

A: [0, 2, 4, 6, 8, 10, 12, 14, 16, 18] B: [0, 3, 6, 9, 12, 15, 18] C: [0, 6, 12, 18]

Python is even so flexible, that it allows to nest list comprehensions in list comprehensions:

# Create the lists
A = [a for a in range(5) if a % 2 == 0]
B = [b for b in range(5) if b % 3 == 0]
C = [a + b for a in A for b in B]

# Print out the result
print('A: ', A)
print('B: ', B)
print('C: ', C)

A: [0, 2, 4] B: [0, 3] C: [0, 3, 2, 5, 4, 7]

Different Data Types

List comprehensions are not only good for numbers, but also other data types can be used. Take for example this list of strings (consisting of verbs):

["work", "eat", "sleep", "repeat"]

verbs = ["work", "eat", "sleep", "repeat"]
verbs_with_ing = [verb + "ing" for verb in verbs]
print(verbs_with_ing)

This results into the expected list:

['working', 'eating', 'sleeping', 'repeating']

To make it even more complex, you could also use a lambda function to convert a verb to its “ing” form:

verbs = ["work", "eat", "sleep", "repeat"]
verb_converter = lambda verb: verb + "ing"
verbs_with_ing = [verb_converter(verb) for verb in verbs]
print(verbs_with_ing)

The point here is that any transformation on the elements can be applied. That makes the list comprehension in Python so powerful.

Python Set Comprehension

Instead of lists, sets can also be used. With sets, lookup is faster and sets have the property that all its elements are unique. One perfect example for Python set comprehension are the prime numbers! Take a look at this Python code:

# Construct a list of integers which are not prime (which are the product of two integers)
no_primes = {a * multiplier for multiplier in range(2, 100) for a in range(2, 100)}
# Since 1 is not a prime number we have to add it to this list
no_primes.add(1)
# Now construct a list of primes out of this list
primes = {p for p in range(1, 100) if p not in no_primes}
# Show the result
print(primes)

This results into the following:

{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97}

Great, this works well!

Python Dictionary Comprehension

Mathematically speaking, a dictionary (dict) is a mapping for which the keys have unique values. In that sense, a dict is a function (since one key only has one value). It is quite easy to define the function

and make a dict comprehension in Python using these lines of code:

# Create a set comprehension for all x in 0...9 for the function f(x)=x^2
f = {x: x**2 for x in range(10)}

# Display the result
print(f)

This results into the following output:

{0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}

Map/Filter Lambda functions

In Python, the map or filter function can be used instead of list/set/dict comprehensions. But, in my point of view this is not good for the readability of your code. Compare the map function to the list comprehension in the following code snippet:

# Compute the lists
U = map(lambda x: x ** 2, [i for i in range(5)])
V = [i ** 2 for i in range(5)]

# Show the results
print('U: ', list(U))
print('V: ', V)

The two lines of code do exactly the same. But what do you think is easier to read? If you have any opinion about it, please let us know in the comments.

Just Keep it Simple

Try to keep things simple. If it gets to complex to write a list, set or dict comprehension, try to use appropriate “if” and “for” statements. Also try to introduce the “map” and “filter” and even lambda functions in that case. It really depends on the use case to go for either list comprehensions or multi-line code statements. If you are confused, please let us know in the comment section!

Conclusions (TL;DR)

It is easy to use the list comprehensions in Python (or Python set or dictionary comprehensions). Besides lists, sets and dictionaries, it is also possible to use different data types like strings, tuples or even lambda functions. The mathematical concepts of lists, sets and functions are found in native Python language. Therefore, Python is a great language for implementing mathematical models since it feel intuitively to implement these ideas. In some cases, it is possible to write the same expression using “if” and “for” statements. If your code gets too complex with list comprehensions, try to make it more readable.