Comprehensions in depth

Part of the free LogicWiz Python course, Module 12: Iteration and the standard toolkit. The examples and practice run in the app.

Shankar bhaiya wants the whole menu with GST added, the list of dishes under ₹20, and a pass/fail list for his son's class test. Meher writes all three in three lines.

You'll learn to

• Write list, dict and set comprehensions, with and without a filter
• Tell a filter `if` (at the end) from a conditional expression `if-else` (at the front)
• Read and write nested comprehensions in the right order
• Know when a normal loop reads better

Learn

Read it like a sentence

[m * 2 for m in marks if m >= 33] reads: give me m * 2, for each m in marks, if m >= 33.

Three parts:

1. what to keep: the expression at the front
2. where it comes from: the for
3. which ones (optional): the if at the end

The result is a brand-new list. The original is not touched.

Filter vs choose

marks = [45, 82, 28, 91, 67]

passed = [m for m in marks if m >= 33]
print(passed)

labels = ["pass" if m >= 33 else "fail" for m in marks]
print(labels)

Where the if goes

Dict and set comprehensions

menu = {"samosa": 15, "chai": 10, "vada pav": 20}

with_gst = {item: round(p * 1.05, 2) for item, p in menu.items()}
print(with_gst)

cheap = {item for item, p in menu.items() if p < 20}
print(sorted(cheap))

word_lengths = {len(w) for w in ["chai", "dosa", "idli", "vada pav"]}
print(word_lengths)

Curly brackets with key: value make a dict; curly brackets with just a value make a set. A set has no fixed order, so we print sorted(cheap).

Nested: same order as the loops

A comprehension with two fors reads in the same order as nested loops:

flat = []
for row in grid:
    for x in row:
        flat.append(x)

becomes [x for row in grid for x in row].

A comprehension inside another one builds a list of lists: [[r * c for c in range(1, 4)] for r in range(1, 4)] is a 3×3 table.

Flatten, table, transpose

grid = [[1, 2, 3], [4, 5, 6]]

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

table = [[r * c for c in range(1, 4)] for r in range(1, 4)]
print(table)

transposed = [[row[i] for row in grid] for i in range(3)]
print(transposed)

Use a loop instead when:

• you need more than two fors or a complicated condition
• you're doing something for its side effect ([print(x) for x in items] builds a useless list of Nones)
• you need try/except or several statements per item

A comprehension should read in one breath.

AKTU, VTU and Anna University papers ask "write a list comprehension for the squares of even numbers from 1 to 20". Interviewers ask "invert a dictionary in one line": {v: k for k, v in d.items()}.

In one line

[expr for x in it if cond]: filter if at the end, a if c else b at the front; {k: v ...} makes a dict, {v ...} a set; nested fors read like nested loops.

Interview corner

What is a list comprehension, and is it faster than a loop?

It is a compact expression that builds a new list: [x * x for x in nums if x > 0]. It is usually a bit faster than an equivalent for loop with .append() because the loop runs in optimised bytecode without looking up append each time. Choose it for clarity, not speed; use a loop when the logic needs statements.

How do you invert a dictionary in one line? What if values repeat?

{v: k for k, v in d.items()}. Values must be hashable to become keys. If several keys share a value, the last one wins; to keep all of them, group instead: inv = defaultdict(list) and inv[v].append(k).

What is the difference between [x for x in it] and (x for x in it)?

Square brackets build the whole list in memory at once. Round brackets make a generator expression that produces items one by one on demand, using almost no memory, and can be looped over only once. Use the generator when you only need to loop or feed sum/max/any.