Skip to main content

List, dict & set comprehensions

Level 2: Level 2: Idiomseasy10 mincomprehensionslistsdictsfiltering

Transform and filter collections in one readable expression.

One expression that says exactly what you want

Reach for a comprehension whenever you are building a new collection out of an existing one by transforming or filtering it. That is most of the collection code you will ever write: pull one field out of a list of API records, drop the rows that fail a validation check, build a lookup table keyed by id. The comprehension puts the result on one line, so a reviewer reads what you are producing instead of tracing an accumulator through a loop body.

The mental model

A list comprehension is a loop-and-collect fused into a single expression.

Same computation, two forms
Explicit loop
squares = []
for n in nums:
    squares.append(n * n)
Comprehension
squares = [n * n for n in nums]
  • output n * n
  • iterate for n in nums
The same loop-and-collect, fused into one expression: the output expression comes first, the loop reads left to right.

Python still runs the loop; it just builds the list for you. Read it left to right as "this expression, for each item in the source."

nums = [1, 2, 3, 4, 5]
squares = [n * n for n in nums]
print(squares)   # [1, 4, 9, 16, 25]

That is the exact shape the Apply exercise wants: take each n, square it, collect the results.

Filter with a trailing if

Add if <condition> after the loop to keep only the items that pass:

Same computation, two forms
Explicit loop
evens = []
for n in nums:
    if n % 2 == 0:
        evens.append(n)
Comprehension
evens = [n for n in nums if n % 2 == 0]
  • output n
  • iterate for n in nums
  • filter if n % 2 == 0
A trailing if maps to the nested if in the loop: it skips items instead of transforming them, so the result gets shorter.
evens = [n for n in nums if n % 2 == 0]
print(evens)     # [2, 4]

Same shape, different brackets

Swap [ ] for { } to build a dict or a set. A dict comprehension needs a key: value pair:

squares_map = {n: n * n for n in nums}
print(squares_map)   # {1: 1, 2: 4, 3: 9, 4: 16, 5: 25}
distinct = {n % 3 for n in nums}   # {0, 1, 2}

The {word: len(word) for word in words} form is precisely what the Practice exercise builds.

Check yourself
With nums = [1, 2, 3, 4, 5], what does [n if n > 2 else 0 for n in nums] produce?

Pitfall: two different if positions

if at the end filters. if ... else at the front is a conditional expression that transforms every item and drops nothing:

[n for n in nums if n > 2]         # [3, 4, 5]        -> filter, shorter
[n if n > 2 else 0 for n in nums]  # [0, 0, 3, 4, 5]  -> map, same length

Interns mix these up constantly. If your output got shorter, you filtered. If it stayed the same length, you mapped.

Check yourself
What is {len(w): w for w in ['hi', 'by', 'ok']}?

One more, for dicts: duplicate keys silently overwrite, and the last one wins. {len(w): w for w in ["hi", "by", "ok"]} keeps only {2: "ok"} because all three keys are 2. In the Practice exercise the words are the keys, so you are safe, but flip the mapping and you will quietly lose data.

Interview nuance: a comprehension is eager. It runs to completion and materializes the whole collection in memory, so it costs O(n) time and O(n) space. That is fine for thousands of items and wasteful for a billion-row scan. The lazy cousin is the generator expression (n * n for n in nums), the same syntax with ( ), which yields one value at a time in O(1) extra space. When an interviewer asks what you would run over a huge stream, "a generator, because the full list would not fit in memory" is the answer they are listening for.

Check yourself

You are choosing between [x for x in src] and (x for x in src). Sort each requirement into the tool that satisfies it.

You loop over the result twice
You total one number over a 10 GB log file
You call len() on the result
You read the result with [0]
You stop at the first match and abandon the rest
Worked example (Python)
nums = [1, 2, 3, 4, 5]
print([n * n for n in nums])              # [1, 4, 9, 16, 25]
print([n for n in nums if n % 2 == 0])    # [2, 4]
print({n: n * n for n in nums})           # {1: 1, 2: 4, ...}

Apply

Your turn

The task this lesson builds to.

Implement squares(nums): return a new list with each number in nums squared.

For [1, 2, 3] return [1, 4, 9]. Use a list comprehension.

3 hints and 4 automated checks are waiting in the workspace.

Practice

Make it stick

A second problem on the same idea, so it survives past today.

Implement lengths(words): return a dict mapping each word to its length.

For ["hi", "abc"] return {"hi": 2, "abc": 3}. Use a dict comprehension.

3 hints and 4 automated checks are waiting in the workspace.