Bytebook Compiler

Data Structures in Python

Comprehensions

Build lists, sets and dictionaries in one readable line, add filters and conditions, nest them, and know when a plain loop is clearer.

3 min read5 runnable examples4-question quiz

A comprehension builds a new collection from an existing one in a single expression:

Python
[expression for item in iterable if condition]
Python
nums = [1, 2, 3, 4, 5, 6, 7, 8]
squares = [n * n for n in nums]
evens = [n for n in nums if n % 2 == 0]
labels = ["even" if n % 2 == 0 else "odd" for n in nums]     # if-else goes BEFORE for
print(squares)
print(evens)
print(labels)
Output
[1, 4, 9, 16, 25, 36, 49, 64]
[2, 4, 6, 8]
['odd', 'even', 'odd', 'even', 'odd', 'even', 'odd', 'even']

The equivalent loop for evens is four lines:

Python
evens = []
for n in nums:
    if n % 2 == 0:
        evens.append(n)

Set and dictionary comprehensions

Python
words = ["apple", "banana", "avocado", "cherry", "blueberry"]
first_letters = {w[0] for w in words}                 # set: unique
lengths = {w: len(w) for w in words}                  # dict
long_words = {w: n for w, n in lengths.items() if n > 6}
print(sorted(first_letters))
print(lengths)
print(long_words)
Output
['a', 'b', 'c']
{'apple': 5, 'banana': 6, 'avocado': 7, 'cherry': 6, 'blueberry': 9}
{'avocado': 7, 'blueberry': 9}

Nested comprehensions

Read the for clauses left to right, exactly as nested loops:

Python
pairs = [(x, y) for x in range(1, 4) for y in range(1, 4) if x != y]
print(pairs)

matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = [v for row in matrix for v in row]
table = [[r * c for c in range(1, 5)] for r in range(1, 4)]
print(flat)
for row in table:
    print(row)
Output
[(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]
[1, 2, 3, 4, 5, 6, 7, 8, 9]
[1, 2, 3, 4]
[2, 4, 6, 8]
[3, 6, 9, 12]

Generator expressions

Use round brackets to get a generator that produces values on demand instead of building a whole list in memory. Perfect inside sum, max, any and all:

Python
print(sum(n * n for n in range(1, 101)))                 # no list created
print(any(c.isdigit() for c in "pass1word"))
print(all(len(w) > 2 for w in ["cat", "dog", "ox"]))
print(max((len(w), w) for w in ["pear", "fig", "banana"]))
Output
338350
True
False
(6, 'banana')

Real examples

Python
marks = {"Asha": 91, "Ravi": 38, "Meera": 85, "Kiran": 29}
passed = [name for name, m in marks.items() if m >= 40]
grade = {name: ("A" if m >= 85 else "B" if m >= 40 else "F") for name, m in marks.items()}
primes = [n for n in range(2, 50) if all(n % d for d in range(2, int(n ** 0.5) + 1))]
print(passed)
print(grade)
print(primes)
Output
['Asha', 'Meera']
{'Asha': 'A', 'Ravi': 'F', 'Meera': 'A', 'Kiran': 'F'}
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

Test yourself

4 questions · pick an answer to check it instantly

  1. 1

    What does [x * 2 for x in range(4) if x % 2] produce?

  2. 2

    Where does an if … else go in a comprehension that keeps every item?

  3. 3

    What is {c for c in "banana"}?

  4. 4

    Why use sum(x * x for x in data) instead of sum([x * x for x in data])?