Comprehensions

July 2021

Table of Contents

Motivating Comprehensions

Regular for loops

# cubing a list of integers
integers = [-4, -3, -2, -1, 0, 1, 2, 3, 4]
cubes = []
for i in integers:
    cubes.append(i ** 3)

print(cubes)  # [-64, -27, -8, -1, 0, 1, 8, 27, 64]

# adding a conditional (use only positive bases)
pos_cubes = []
for i in integers:
    if i > 0:
        pos_cubes.append(i ** 3)

print(pos_cubes)  # [1, 8, 27, 64]

Built-in `map()` and `filter()`

# map
map_cubes = map(lambda x: x ** 3, integers)
map_cubes = list(map_cubes)
print(map_cubes)  # [-64, -27, -8, -1, 0, 1, 8, 27, 64]
assert map_cubes == cubes

# map & filter
map_pos_cubes = map(lambda x: x ** 3, filter(lambda x: x > 0, integers))
map_pos_cubes = list(map_pos_cubes)
print(map_pos_cubes)  # [1, 8, 27, 64]
assert map_pos_cubes == pos_cubes

Note: even with simple lambdas, the expression is not very readable and requires too much boilerplate.

Enter List Comprehensions

# appreciate the readability
lc_cubes = [i ** 3 for i in integers]
print(lc_cubes)  # [-64, -27, -8, -1, 0, 1, 8, 27, 64]
assert map_cubes == cubes == lc_cubes

List Comprehension with Conditionals

lc_pos_cubes = [i ** 3 for i in integers if i > 0]
print(lc_pos_cubes)  # [1, 8, 27, 64]
assert map_pos_cubes == pos_cubes == lc_pos_cubes

Two-dimensional List Comprehension

# motivation
matrix = [[9, 8, 7], [6, 5, 4], [3, 2, 1, 0]]
unrolled = []
for row in matrix:
    for scalar in row:
        unrolled.append(scalar)
print(unrolled)  # [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

# list comprehension approach
lc_unrolled = [scalar for row in matrix for scalar in row]
print(lc_unrolled)  # [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
assert unrolled == lc_unrolled

# unroll only even entries
lc_e_unrolled = [scalar for row in matrix for scalar in row if scalar % 2 == 0]
print(lc_e_unrolled)  # [8, 6, 4, 2, 0]

Note: list comprehensions allow for succinct and readable expressions, up to a point. Even though it is syntactically possible to add more than two for sub-expressions and one extra conditional per for sub-expression, at that point readability goes down substantially. For such cases a traditional for loop is preferable.

Dictionary Comprehensions

From an Iterable

days = ["måndag", "tisdag", "onsdag", "torsdag", "fredag", "lördag", "söndag"]
dict_days = {k: v[:3] for k, v in enumerate(days)}
print(dict_days)
# {0: 'mån', 1: 'tis', 2: 'ons', 3: 'tor', 4: 'fre', 5: 'lör', 6: 'sön'}

Build On the Fly

# dictionary with x square
dict_squares = {x: x ** 2 for x in range(7)}
print(dict_squares)  # {0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36}

Set Comprehensions

# only cube even numbers and return them in a set
items = [3, 3, 3, 2, 2, 2, 1, 1, 1, 4, 4, 4]
cubed_set_evens = {x ** 3 for x in items if x % 2 == 0}
print(cubed_set_evens)  # {8, 64}

Rafael Marino