List Comprehension in Python

In the last article you learned to generate lists with for loops. In this article you’ll learn to generate lists using list comprehensions.

Multiple lines of for loop and if statements can be combined into one line of code.

Without list comprehensions, we’ll find the squares of numbers like this:

squares = []
for num in range(1,11):
    squares.append(num**2)
print(squares)

With list comprehension:

squares = [num**2 for num in range(1,11)]
print(squares)

The expression for generating elements of the new list is put first and the iterable to be considered is put after the for keyword.

The previous code takes all elements in the range object, squares it and appends it to the new list. It is not required to convert the range object to a list.

if Statements in List Comprehensions

if statements can be included in list comprehensions to test condition(s).

The even numbers code from the previous article can be combined into a list comprehension

even_numbers = [num for num in range(1,11) if num%2 == 0]
print(even_numbers)

Else statements can also be included in list comprehensions.

Say you are a school teacher. Your students have just performed poorly on an exam. Their scores have been represented with a list.

exam_scores = [33, 72, 88, 65, 30, 35, 66]

Only few of them performed well. You don’t want to record these poor scores so you’ve decided to add few marks to their scores. This approach can push the highest score beyond the maximum obtainable score.

You can add marks for only students who are below the pass mark, 40. The if-else chain can be used to test the condition for adding scores in our list comprehension.

exam_scores = [33, 72, 88, 65, 30, 35, 66]
print(f'initial_scores: {exam_scores}')

upgraded_scores = [(score + 10) if (score < 40) else score for score in exam_scores]
print(f'upgraded_scores: {upgraded_scores}')

The above code adds 10 marks to scores less than 40 and does nothing to scores not less than 40.

However, there are better ways to normalize these test scores. This is just an illustration of how if-else chains can be used in list comprehensions.

Multiple if-else Chains in List Comprehension

There can be multiple if-else chains in list comprehensions.

You’re building a machine learning model to predict which of your customers would cancel their subscription in the following month.

There’s a feature that represents how long a subscriber has been active on this service. The values contained in this feature can be grouped on quarterly basis.

Let’s say the feature is represented with a list as shown below

tenure = [3, 13, 8, 3, 11, 1,15, 7, 10, 2, 6, 4]

This can be grouped into 4 categories:

A – 0 – 4 months

B – 4 – 8 months

C – 8 – 12 months

D – 12 months or more

tenure = [3, 13, 8, 3, 11, 1,15, 7, 10, 2, 6, 4]

tenure_transformed = ['A' if x <= 4
                     else 'B' if 4<x<=8
                     else 'C' if 8<x<=12
                     else 'D' for x in tenure]

print(tenure_transformed)

We now have 4 categories for the tenure feature.

Dictionary Comprehension

Just like list comprehension, it s possible to have dictionary comprehension.

Let’s create a dictionary to include x values and y values for a quadratic expression

points = {x:(x**2 - 25) for x in range(-10,10)}
print(points)

Each element in the range object is taken and plugged into the quadratic expression. The elements in the range object are the keys of the dictionary while the resulting values from the quadratic expression are the values of the dictionary.

The excerpt from the output shows that the new dictionary was created successfully.

We can go ahead to plot these points on a simple graph.