List Comprehensions

A list comprehension builds a new list from an existing sequence in a single line. It replaces a common pattern seen on the For Loops page: create an empty list, loop over something, and append to it on each pass.

The Python documentation has a guide on List Comprehensions.

Syntax

A list comprehension follows this pattern:

[expression for item in iterable]

item takes on each value in iterable in turn, expression is evaluated using that value, and the results are collected into a new list.

# The "manual" way, using a for loop
squares = []
for i in range(10):
    squares.append(i ** 2)
print(squares)

# The same thing with a list comprehension
squares = [i ** 2 for i in range(10)]
print(squares)

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

Both loops above produce the same list. The comprehension just does it in one line instead of three.

Filtering with a Condition

Adding if condition after the for clause skips any item for which the condition is false:

evens = [i for i in range(20) if i % 2 == 0]
print(evens)

# Same thing the "manual" way, for comparison:
evens_manual = []
for i in range(20):
    if i % 2 == 0:
        evens_manual.append(i)
print(evens_manual == evens)

[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
True

Only values of i where i % 2 == 0 is true make it into evens. Every other value is skipped entirely, not included as some placeholder.

Comprehension with an if/else Expression

A different use of if inside a comprehension is a conditional expression, written before the for clause, which chooses between two values rather than filtering:

altitudes = [400, 20000, 35786, 100000]

labels = ["high" if h > 10000 else "low" for h in altitudes]
print(labels)

['low', 'high', 'high', 'high']

Every altitude produces a label, either "high" or "low", so the output list is always the same length as the input. This is different from the filtering form above: if before the for chooses a value for every item, while if after the for decides whether an item is included at all.

Looping Over Two Sequences with zip()

for in a comprehension can iterate over zip() of two sequences, pairing up corresponding elements from each:

names = ["booster", "second stage", "fairing"]
masses = [25000, 4500, 1900]

descriptions = [f"{name}: {mass} kg" for name, mass in zip(names, masses)]
print(descriptions)

['booster: 25000 kg', 'second stage: 4500 kg', 'fairing: 1900 kg']

zip(names, masses) pairs each name with the mass at the same position, and the comprehension builds one description string per pair.

Calling a Function Inside a Comprehension

The expression in a comprehension can be a function call, applying that function to every item in the sequence:

def classify_signal(rssi_dbm):
    return "Strong" if rssi_dbm > -70 else "Weak"

readings = [-52, -68, -74, -90, -61]
strengths = [classify_signal(r) for r in readings]
print(strengths)

['Strong', 'Strong', 'Weak', 'Weak', 'Strong']

classify_signal is called once per reading, and the comprehension collects the results into strengths.

Nested Comprehensions

A comprehension can be nested inside another, which is one way to build a 2D grid:

identity = [[1 if i == j else 0 for j in range(3)] for i in range(3)]
print(identity)

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

The outer comprehension builds one row per value of i, and the inner comprehension builds that row’s three columns. Nested comprehensions get hard to read quickly. If you find yourself nesting more than two levels, a regular for loop is often clearer.

A comprehension can also use two for clauses side by side, rather than nested, to flatten a list of lists into a single flat list:

grid = [[0, 0, 0], [0, 1, 0], [0, 0, 0]]

flat = [value for row in grid for value in row]
print(flat)

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

This reads as “for each row in grid, for each value in that row, keep value.” That’s the same order the two for clauses would appear in an equivalent nested for loop.

Reading Questions

What problem does a list comprehension solve compared to writing a for loop that appends to an empty list?
Write a list comprehension that produces the cubes of the numbers from 0 to 9.
In [i for i in range(10) if i % 3 == 0], what does the if clause do to the resulting list?
What is the difference between if used as a filter (after for) and if used as a conditional expression (before for) in a list comprehension?
Can a list comprehension call a function as its expression? Give an example.
What does zip(names, masses) produce when used inside a comprehension’s for clause?
How would you flatten [[1, 2], [3, 4]] into [1, 2, 3, 4] using a comprehension?