Table of Contents

Introduction

List comprehensions provide a concise way to create lists. It is typically used to create new lists where each element is the result of some operations applied to each member of another sequence.

List Comprehension

Assume we want to create a list of squares of the numbers from 1 to 10 using a for loop.

Example

for loop to create a list.

1squares = []
2for i in range(1,11):
3    squares.append(i**2)
4squares

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

List comprehension offers a more compact syntax when you want to create a new list based on the values of an existing list. Using a list comprehension, we can reproduce the above results as follows.

Example

List comprehension to create a list.

1[x**2 for x in range(1,11)]

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

More generally, a list comprehension consists of square brackets containing an expression followed by a for clause, with optional 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.

Syntax

List Comprehension with an if clause.

1[expr for val in collection if condition]

The above is equivalent to the following for loop enclosing an if clause.

1result = []
2    for val in collection:
3        if condition:
4            result.append(expr)

As an example, imagine we are given a list of names from which we wish to extract only those that start with “J” and capitalize them. A for loop method will look like this:

Example

Filtering a list with a for loop.

1collection = ["Jane", "Peter", "Tom", "Julian", "Wendy", "Violet"]
2result = []
3for val in collection:
4    if val[0]== "J":
5        result.append(val.upper())
6result

['JANE', 'JULIAN']

On the other hand, a list comprehension will look like this.

Example

Filtering a list with a list comprehension.

1[val.upper() for val in collection if val[0]=="J"]

['JANE', 'JULIAN']

As another example, assume we want to create a list of perfect squares less than 200 (starting from 1). Using a while loop, we arrive at the following:

Example

Creating a list of perfect squares with a while loop.

1squares = []
2i = 1
3while i**2 < 200:
4    squares.append(i**2)
5    i +=1
6squares

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196]

On the other hand, using a list comprehension, we can write:

Example

Creating a list of perfect squares with a list comprehension.

1[x**2 for x in range(1,999) if (x**2) < 200]

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196]

In the above code, we are creating a new list using the expression x**2 where x comes from a (purposefully large) range object, subjected to the condition $x^2 < 200$.

We can also include conditions in the expression to manipulate the outcome. In the following example, given a list of names, those that start with “J” are set to 1 and the rest are set to 0. A simple sum() function will then return the number of names that start with “J”.

Example

Including a condition in the list comprehension expression.

1names = ["Jane", "Peter", "Tom", "Julian", "Wendy", "Violet"]
2io = [(1 if x[0]=="J" else 0 ) for x in names]
3print(io)
4print(sum(io))

[1, 0, 0, 1, 0, 0]
2

The expression may also be a Boolean expression. Let’s rework the above example with the expression now replaced by a Boolean expression.

1names = ["Jane", "Peter", "Tom", "Julian", "Wendy", "Violet"]
2io = [x[0]=="J" for x in names]
3print(io)
4print(sum(io))

[True, False, False, True, False, False]
2

ADVERTISEMENT

Nested List Comprehension

Imagine now we want to create a matrix by using 2 for loops, one nested within the other.

Example

Create a matrix by using nested for loops.

1matrix = []
2for i in range(3):   # row
3    matrix.append([])
4    #print(matrix)
5    for j in range(5): # col
6        matrix[i].append(j)
7print(*matrix, sep='\n')

[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]

Let’s reproduce the above results using a nested list comprehension which we will elaborate in 3 steps so that the reader can better understand.

Step 1| Create simple list comprehension.

1[i for i in range(3)] #Step 1: simple list comprehension

[0, 1, 2]

Step 2| Add an empty list as the expression.

1[[] for i in range(3)] # Step 2: add an empty list as the expression

[[], [], []]

Step 3| Write a for loop in the empty list.

1[[j for j in range(5)] for i in range(3)] # Step 3: populate the empty lists

[[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]
Tip

To print a nested list in matrix format. Use print(*list, sep = '\n').

1mat = [[j for j in range(5)] for i in range(3)] # from step 3
2print(*mat, sep='\n')

[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]

Alternatively, we can also use a for loop to print the matrix.

1for row in mat:
2    print(f'{row}')

[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]

As a final example, let’s create a 4 X 4 matrix of random integers from 10 to 99.

Example

Creating a $4 \times 4$ matrix of random integers.

1import random as rd
2A= [[j for j in list(rd.sample(range(10,100), 4))] for i in range(4)]
3# print out the result
4for row in A:
5    print(f'{row}')

[35, 47, 42, 91]
[89, 43, 60, 85]
[64, 86, 85, 62]
[86, 41, 31, 14]

Example

Count how many of the numbers generated are greater than 50.

1total=0
2for row2 in [[j>50 for j in row1]  for row1 in A]:
3    total += sum(row2)
4total

9

A more compact way.

1sum([sum(j>50 for j in row)  for row in A])

9

Dictionary Comprehension

We also have dictionary comprehensions. In the following, we use a dict comprehension to group the scores of 6 students into either “PASS” or “FAIL”.

Example

Dictionary comprehension.

1scores = {"Carol":49, "Sammy":85, "Derrick":36, "Juan":81, "Wendy":79, "Tom":52}
2{k: "PASS" if v > 50  else "FAIL" for (k, v) in scores.items()}

{'Carol': 'FAIL',
 'Sammy': 'PASS',
 'Derrick': 'FAIL',
 'Juan': 'PASS',
 'Wendy': 'PASS',
 'Tom': 'PASS'}

We may also include more than 2 categories in the expression.

Example

Dictionary comprehension with multiple categories.

1scores1 = {"Carol":68, "Sammy":85, "Derrick":62, "Juan":49, "Wendy":79, "Tom":86}
2{k:
3 "A" if v >85 else
4 "B" if v > 75 else
5 "C" if v >65 else
6 "D"
7 for (k, v) in scores1.items()}

{'Carol': 'C',
 'Sammy': 'B',
 'Derrick': 'D',
 'Juan': 'D',
 'Wendy': 'B',
 'Tom': 'A'}

Set Comprehension

Recall that a set has only unique elements. Therefore, the set comprehension always produces a set with unique elements even if the input list has duplicate elements.

Example

Set comprehension.

1currencies = ["sgd", "usd", "hkd", "eur", "aud", "inr", "rmb", "sgd", "eur"]
2{x.upper() for x in currencies}

{'AUD', 'EUR', 'HKD', 'INR', 'RMB', 'SGD', 'USD'}

Example

Set comprehension with an if clause.

1{x.upper() for x in currencies if x[2]=="d"} # only currencies that end with a "d".

{'AUD', 'HKD', 'SGD', 'USD'}