Comprehensions

Set-builder notation come to life

Mathematics has a compact way to describe a collection built from another: the set-builder. The comprehension is that notation, typed directly:

Read “the set of $x^2$, for each $x$ in this source”, and the Python is the same sentence reversed into code:

# Regular loop
squares = []
for x in range(6):
    squares.append(x ** 2)

# Comprehension
squares = [x ** 2 for x in range(6)]
# [0, 1, 4, 9, 16, 25]

The loop spells out three moves, start empty, append, repeat; the comprehension states the whole collection in one line that mirrors the set-builder’s anatomy: the expression up front, the ranged variable behind it.

Filtering with conditions

Set-builder notation also carries membership tests. $\{w \in words \mid |w| > 2\}$ becomes a trailing if:

evens = [x for x in range(10) if x % 2 == 0]
# [0, 2, 4, 6, 8]

long_words = [w.upper() for w in ["hi", "hello", "hey"] if len(w) > 2]
# ['HELLO', 'HEY']

An if at the end is a filter: only the elements that pass it reach the expression. The element travels expression → filter → list, in the order the sentence reads.

If/else as an expression

The if...else you already know is an expression, it produces a value. Sticking one before the for plants it in the build line, choosing per element rather than filtering per element:

labels = ["even" if x % 2 == 0 else "odd" for x in range(5)]
# ['even', 'odd', 'even', 'odd', 'even']

The two positions are a fork with distinct jobs: after the for, the clause votes on elements; before the for, it labels them. One discards, the other transforms.

Dict comprehensions

The same shape builds mappings, expression on the left of the colon becomes the key, expression on the right the value:

squares_dict = {x: x**2 for x in range(6)}
# {0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25}

# Invert a dict
original = {"a": 1, "b": 2}
inverted = {v: k for k, v in original.items()}
# {1: 'a', 2: 'b'}

The inversion is the elegant classic: walk items(), and swap which half of each pair becomes the key.

Set comprehensions

Braces with a comprehension yield a set, uniqueness applied automatically:

lengths = {len(word) for word in ["hello", "hi", "hey"]}
# {2, 3, 5}  (unique lengths)

Three lengths collapse to a set of values, dropping the duplicate as a set must.

Nested comprehensions: flattener

A matrix is a list of rows, and flattening it is two loops in one expression, read the for clauses left to right, outer first:

matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = [num for row in matrix for num in row]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

Each for unwraps one level: row iterates the outer list, num iterates each row, and the collection’s order follows the loops exactly.

A worked example: three lines from the set-builder

The lesson’s three moves, build, filter, label, one line each:

squares = [x ** 2 for x in range(2, 9)]
# [4, 9, 16, 25, 36, 49, 64]

numbers = [x for x in range(1, 11) if x % 3 == 0]
# [3, 6, 9]

labels = ["even" if x % 2 == 0 else "odd" for x in numbers]
# ['odd', 'even', 'odd']

The first is $\{x^2 \mid x \in [2, 9)\}$ typed directly; the second filters the divisors of $3$; the third labels each survivor. What the set-builder states in one breath, the comprehension spells in one line.

When NOT to use comprehensions

• When the logic turns knotty, a regular for loop earns its readability.
• When the body needs try/except, comprehensions have no room for it.
• When side effects matter, printing or writing files should be deliberate statements, not silent expressions.
• Putting if before the for labels instead of filtering. [x if x % 2 == 0 else 'odd' for x in ...] keeps every element, merely relabeled; only an if after the for discards. Slot it wrong and the rejects silently stay.