List and Dictionary Comprehension in Python

March 10, 2022

List and Dictionary Comprehension in Python

There's even a Pythonic way of writing code - which is no less than to write code that just doesn't get the syntax right but that follows the conventions of the Python community and uses the language in the way is intended to be used. (more about on Python's convention on PEP8 style guide )

Following several suggestions on how to write code in a more Pythonic way, we found the use of List comprehension and Dict comprehension.

List Comprehension

What is a list comprehension ? It's a compact way to process all or part of the elements in a sequence and return a list with the results.

Let's suppose we want to transform a string using a for loop:

p_list = []
for letter in 'Pink Floyd':
    p_list.append(letter)

print(p_list)

We would get:

['P', 'i', 'n', 'k', ' ', 'F', 'l', 'o', 'y', 'd']

Now, let's do the same using list comprehension.
p_list = [letter for letter in 'Pink Floyd']

print(p_list)

Same result.
['P', 'i', 'n', 'k', ' ', 'F', 'l', 'o', 'y', 'd']

The syntax of List Comprehension is as follows:
[expression for item in list]



It consists of brackets containing an expression followed by a "for" clause, then zero or more "for" or "if" clauses. The result will be a new list resulting from evaluating the expression in the context of the "for" and "if" clauses which follow it.
If you have a keen eye, you will recognize 'Pink Floyd' is not a list but a string. This is the power of List comprehension: it can identify when it receives a string or a tuple and work on it like a list. The idea is to write more concise and readable code. 
Next: you can modify the values of your list
p_list = [x**2 for x in range(10)]

print(p_list)
This code will give you:
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81] 

Or even add a conditional statement (a real ecstasy for one-line lovers):
p_list = [x for x in range(20) if x % 2 == 0]

print(p_list)
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

We can also have nested IF statements
m_list = ['Genesis', 'Pink Floyd', 'Yes', 'Muse']
p_list = [x for x in m_list if x[0] == 'G' if len(x) == 7]

print(p_list)
You will get:
['Genesis']

Nested List Comprehension
Consider the following example of a 4x4 matrix implemented as a list of 4 lists of length 4.
matrix = [
    ['Genesis', 'Pink Floyd', 'Yes', 'Muse'],
    ['Phil', 'David', 'Rick', 'Matt'],
    ['Peter', 'Roger', 'Jon', 'Chris'],
    ['Tony', 'Richard', 'Chris', 'Dominic']
]
The following list comprehension will transpose rows and columns:
result = [[row[i] for row in matrix] for i in range(4)]

print(result)
You will get:
[['Genesis', 'Phil', 'Peter', 'Tony'], ['Pink Floyd', 'David', 'Roger', 'Richard'], ['Yes', 'Rick', 'Jon', 'Chris'], ['Muse', 'Matt', 'Chris', 'Dominic']]

If you wanted to do the same using a normal for loop, you would have to:
transposed = []
matrix = [
    ['Genesis', 'Pink Floyd', 'Yes', 'Muse'],
    ['Phil', 'David', 'Rick', 'Matt'],
    ['Peter', 'Roger', 'Jon', 'Chris'],
    ['Tony', 'Richard', 'Chris', 'Dominic']
]
for i in range(4):
    transposed_row = []
    for row in matrix:
        transposed_row.append(row[i])
    transposed.append(transposed_row)

print(transposed)
[['Genesis', 'Phil', 'Peter', 'Tony'], ['Pink Floyd', 'David', 'Roger', 'Richard'], ['Yes', 'Rick', 'Jon', 'Chris'], ['Muse', 'Matt', 'Chris', 'Dominic']]

You can see the difference between the for loop and the list comprehension. Again: the idea is to do less code, more readable and easier to maintain. List comprehension is generally more compact and faster than normal functions and loops for creating lists. 
Dictionary Comprehension
First of all, what is a Python Dictionary? A set of key:value pairs, with the requirement that the keys are unique. 
A dictionary comprehension takes in a group of values, performs some kind of task on them, and returns the result in the form of a dictionary.
 
The general syntax for a dictionary comprehension looks like this:
{resulting_keys: resulting_values for iteration in series}
The input for a dictionary comprehension does not have to be a dictionary, but the simplest way to explain it would be to start off converting one dictionary to another.
We have the following dictionary:
diction = {'genesis': 1, 'oasis': 2, 'muse': 3, 'yes': 4, 'rem': 5}
And we want to modify all the values (multiplied by 2). This can be done with a simple dictionary comprehension, like this 
new_dict = {key: value * 2 for key, value in diction.items()}
The new dictionary will look like this:
new_dict = {'genesis': 2, 'oasis': 4, 'muse': 6, 'yes': 8, 'rem': 10}
If we compare this syntax from the list comprehension, we can see we're now iterating over two groups in tandem, a group of keys and a group of values. The items() method converts the dictionary into a list of tuples that can be read as (key1, value1), (key2, value2), etc.
Let's say now we also want to modify the keys and the values at the same time, we could:
new_dict = {key.capitalize(): value * 2 for key, value in diction.items()}
We would have:
new_dict = {'Genesis': 2, 'Oasis': 4, 'Muse': 6, 'Yes': 8, 'Rem': 10}
In a way, we can think we're performing two lists comprehensions at once: one for the keys, another one for the values. 
Conditionals within Dictionary Comprehension

As with lists, we can also filter and/or add conditional logic inside a dictionary comprehension. It works the same way as the lists, so:
diction = {'genesis': 1, 'oasis': 2, 'muse': 3, 'yes': 4, 'rem': 5}
new_dict = {key: value for key, value in diction.items() if value > 2}
This will give us only the key-value pairs in which the value is greater than 2:
{'muse': 3, 'yes': 4, 'rem': 5}
We can also filter on the keys. Let's add one more filter:
new_dict = {key: value for key, value in diction.items() if value > 2 if 'mu' in key}
We would get:
{'muse': 3}

J'capote!!
As you can see, we just need to add more "if" to the same line.


Nested Dictionary Comprehension
Like list comprehensions, dictionary comprehensions can be nested within each other and now things start to get interesting. Imagine we have a dict of employees and we want to filter them based on the salary:
diction = {
    'emp_1': {
        'last_name': 'Jones',
        'first_name': 'Mr.',
        'salary': 40000},
    'emp_2': {
        'last_name': 'Manilow',
        'first_name': 'Barry',
        'salary': 75000},    'emp_3': {
        'last_name': 'Eastwood',
        'first_name': 'Clint',
        'salary': 90000}
}
new_dict = {outer_key: outer_val for outer_key, outer_val in diction.items() if outer_val['salary'] > 50000}

Only Barry and Clint would make it to that result:
{'emp_2': {'last_name': 'Manilow', 'first_name': 'Barry', 'salary': 75000}, 'emp_3': {'last_name': 'Eastwood', 'first_name': 'Clint', 'salary': 90000}}

Now let's use the same dict but this time let's use a conditional "if" and "else" to multiply by 2the salary of our fellow employees.  
new_dict = {outer_key: {inner_key: (inner_val * 2 if inner_key == 'salary' else inner_val) for inner_key, inner_val in outer_val.items()} for outer_key, outer_val in my_dict.items()}

Whoa! What just happened there ?  The result: 
 
{'emp_1': {'last_name': 'Jones', 'first_name': 'Mr.', 'salary': 80000}, 'emp_2': {'last_name': 'Manilow', 'first_name': 'Barry', 'salary': 150000}, 'emp_3': {'last_name': 'Eastwood', 'first_name': 'Clint', 'salary': 180000}}
Before your head starts spinning round, let me explain: try to think what the inner dictionary is doing first (multiplying the salary by 2). Remember that the result of that as a whole represents the new set of values for the outer dictionary. It can be thought as: 
new_dict = {outer_key: {outer_val} for outer_key, outer_val in my_dict.items()}
The outer_val is actually its own dictionary, which we are doing a dictionary comprehension on for each inner dictionary. 
That's it! Don't forget: we want to code in a more readable way, easier to understand; if you feel more comfortable doing a classic for loop, go for it. Try always to do it in a pythonic way.