Chapter 9 Iterators, Generators, and Comprehensions

Iteration is central to Python coding interviews and data processing. The key distinction is between an iterable, which can produce an iterator, and an iterator, which produces values one at a time while maintaining state.

9.1 Iterable versus iterator

Lists, tuples, dictionaries, strings, and many pandas objects are iterable. Calling iter() creates an iterator, and next() requests its next value.

values = [10, 20, 30]
iterator = iter(values)

assert next(iterator) == 10
assert next(iterator) == 20
assert next(iterator) == 30

After the final value, another next() raises StopIteration. A for loop handles that exception internally.

iterator = iter([1])
assert next(iterator) == 1

try:
    next(iterator)
except StopIteration:
    exhausted = True

assert exhausted

9.2 Implementing the iterator protocol

A custom iterator implements __iter__() and __next__(). __next__() must raise StopIteration when no values remain.

class Countdown:
    def __init__(self, start):
        self.current = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.current <= 0:
            raise StopIteration

        value = self.current
        self.current -= 1
        return value


assert list(Countdown(3)) == [3, 2, 1]

For simple sequences, a generator is normally shorter and easier to maintain. Custom iterator classes are most useful when iteration requires meaningful state or additional behavior.

9.3 Generators produce values lazily

A function containing yield returns a generator. It pauses after yielding a value and resumes from the same state when the next value is requested.

def nonmissing(values):
    for value in values:
        if value is not None:
            yield value


clean_values = list(nonmissing([3, None, 5, None, 8]))
assert clean_values == [3, 5, 8]

Generators are useful when the full result would be large or when the consumer may stop early. They are single-use: once exhausted, create a new generator to iterate again.

9.4 Generator expressions

A generator expression resembles a list comprehension but uses parentheses. It produces one item at a time instead of building the complete list.

squares = (number ** 2 for number in range(1, 6))

assert sum(squares) == 55

Use a list comprehension when you need the complete reusable result. Use a generator expression when a streaming, single-pass result is sufficient.

9.5 Comprehensions

Comprehensions are concise when the transformation and filter remain easy to read. Do not compress complex branching or side effects into one expression.

even_squares = [number ** 2 for number in range(10) if number % 2 == 0]
square_lookup = {number: number ** 2 for number in range(4)}

assert even_squares == [0, 4, 16, 36, 64]
assert square_lookup == {0: 0, 1: 1, 2: 4, 3: 9}

9.6 `enumerate()` and `zip()`

Use enumerate() instead of manually managing a counter. Use zip() to traverse aligned iterables. By default, zip() stops at the shortest input, so validate lengths when silent truncation would be dangerous.

features = ["age", "income", "tenure"]
positions = {feature: index for index, feature in enumerate(features)}

names = ["Ada", "Grace"]
scores = [95, 98]
paired_scores = list(zip(names, scores))

assert positions["income"] == 1
assert paired_scores == [("Ada", 95), ("Grace", 98)]

In modern Python, zip(..., strict=True) can detect unequal input lengths:

try:
    list(zip([1, 2], ["a"], strict=True))
except ValueError:
    length_mismatch_detected = True

assert length_mismatch_detected

9.7 Complexity and memory reasoning

If an iterable contains \(n\) values, visiting every value is generally \(O(n)\) time. A list comprehension that retains every result also requires \(O(n)\) additional memory. A generator can use \(O(1)\) additional memory when it holds only its current state, although the operations performed for each item still determine total time complexity.

9.8 Interview checklist

Is the object iterable, an iterator, or both?
Will the result be consumed once or reused?
Can lazy evaluation reduce peak memory?
Could zip() silently truncate mismatched inputs?
Is a comprehension still readable?
Does the solution need indices, values, or both?

9.9 Exercises

Write a generator that yields rolling pairs from [1, 2, 3, 4] as (1, 2), (2, 3), and (3, 4).
Convert a loop that appends cleaned strings to a list comprehension.
Explain why converting a large generator immediately to list() removes its peak-memory advantage.