InterviewPrepKit

Home / Learn / Algorithms & Data Structures with Python

Comprehensions and Iterators

By Armin Norouzi

What this lesson is about

You already know how to loop over things and collect results. This lesson covers two ideas that make that work shorter and clearer, and that explain what is really happening when you loop.

The first idea is the comprehension: a compact way to build a list, dictionary, or set from an existing collection in a single line.

The second idea is the iterator protocol: the small set of rules Python follows every time a for loop runs. Once you understand it, loops stop being magic. You will also see two everyday helpers, enumerate and zip, that lean on this protocol.

List comprehensions

A list is an ordered collection of values, written with square brackets: [10, 20, 30]. Suppose you want a new list where every number is doubled. The long way uses an explicit loop:

numbers = [1, 2, 3, 4]
doubled = []
for n in numbers:
    doubled.append(n * 2)
print(doubled)  # -> [2, 4, 6, 8]

Here append adds one item to the end of a list. The pattern is: start with an empty list, loop, and append a transformed value each time. That pattern is so common that Python gives you a one-line form for it, called a list comprehension:

numbers = [1, 2, 3, 4]
doubled = [n * 2 for n in numbers]
print(doubled)  # -> [2, 4, 6, 8]

Read it left to right: “the value n * 2, for each n in numbers.” The part before for is the expression (what to produce each time). The part after for is the loop. It produces exactly the same list as the four-line version.

Adding a filter condition

You can keep only some items by adding an if at the end. This is called the filter condition: an item is included only when the condition is true.

numbers = [1, 2, 3, 4, 5, 6]
evens = [n for n in numbers if n % 2 == 0]
print(evens)  # -> [2, 4, 6]

The % operator gives the remainder after division, so n % 2 == 0 is true when n is even. Items that fail the condition are skipped, not transformed.

You can transform and filter at the same time. Here we square only the even numbers:

numbers = [1, 2, 3, 4, 5, 6]
even_squares = [n * n for n in numbers if n % 2 == 0]
print(even_squares)  # -> [4, 16, 36]

The order inside the brackets is fixed: expression first, then for, then an optional if.

Cost

A comprehension over n items does a constant amount of work per item, so it runs in O(n) time (time grows in proportion to the number of items). It builds a new list holding up to n results, so it uses O(n) space (extra memory in proportion to the number of items). It is not faster in the Big-O sense than the explicit loop; it is the same work written more compactly, with a small constant-factor speedup in practice.

Dict and set comprehensions

The same shape works for two other collections.

A dictionary (or dict) maps keys to values, written with curly braces and colons: {"a": 1, "b": 2}. A dict comprehension produces key: value pairs:

words = ["hi", "hey", "hello"]
lengths = {w: len(w) for w in words}
print(lengths)  # -> {'hi': 2, 'hey': 3, 'hello': 5}

len(w) gives the number of characters in the string w. The expression before for is now w: len(w), a key followed by its value.

A set is an unordered collection with no duplicates, also written with curly braces but with no colons: {1, 2, 3}. A set comprehension builds one and drops repeats automatically:

numbers = [1, 2, 2, 3, 3, 3]
unique_squares = {n * n for n in numbers}
print(unique_squares)  # -> {1, 4, 9}

Even though 2 and 3 appear several times, each squared value is stored once.

Note the ambiguity: {} alone is an empty dict, not an empty set. For an empty set you must write set().

Readability: when to use a comprehension

A comprehension is the right tool when it stays short and reads like a sentence. It is the wrong tool when it gets long or does several things. Compare:

# Clear: one transform, one filter.
active_names = [u["name"] for u in users if u["active"]]

If you find yourself nesting conditions, calling several functions, or wanting a print or error handling inside the brackets, use an explicit loop instead. You cannot put statements like if x: ... else: ... blocks, try, or multiple steps inside a comprehension anyway, and forcing complex logic into one line makes it harder to read, not easier. The goal is clarity. A three-line loop that another reader can follow is better than a clever one-liner they cannot understand.

What “iterable” means

An iterable is any object you can loop over with a for loop: lists, strings, dictionaries, sets, ranges, files, and more. When you write for x in something, that something must be iterable.

Being iterable means the object can hand out its items one at a time, on request. It does not mean all the items exist in memory at once. A range, for example, is iterable but does not store every number; it produces them as needed.

for ch in "hi":
    print(ch)
# -> h
# -> i

A string is iterable, and looping over it yields one character at a time.

Iterators, iter(), and next()

An iterator is the object that actually does the handing-out. It remembers where it is in the sequence and produces the next item each time you ask.

Two built-in functions expose this directly:

  • iter(x) takes an iterable and returns a fresh iterator positioned at the start.
  • next(it) asks an iterator for its next item and advances its position by one.
numbers = [10, 20, 30]
it = iter(numbers)
print(next(it))  # -> 10
print(next(it))  # -> 20
print(next(it))  # -> 30

Each next call moves forward. The iterator holds the position; the original list is unchanged.

StopIteration

When there are no items left, next raises an exception (a signal that something happened that stops normal flow) called StopIteration:

it = iter([1])
print(next(it))  # -> 1
print(next(it))  # raises StopIteration

This is not an error in your logic. It is the agreed-upon way for an iterator to announce “I am finished.” You rarely call next yourself, so you rarely see this exception; the for loop catches it for you, as the next section shows.

The distinction to hold on to: an iterable is something you can get an iterator from (like a list); an iterator is the moving cursor that yields items and eventually raises StopIteration.

How a for loop works under the hood

A for loop is built entirely from iter, next, and StopIteration. When you write:

for x in numbers:
    print(x)

Python effectively does this:

it = iter(numbers)          # 1. get an iterator
while True:
    try:
        x = next(it)        # 2. pull the next item
    except StopIteration:   # 3. no items left -> stop
        break
    print(x)                # 4. run the loop body

Step 1 happens once. Steps 2 through 4 repeat: pull an item, run the body, pull the next. When next raises StopIteration, the loop ends cleanly. That is the whole mechanism. Every for loop you have ever written is this pattern.

The same cycle, drawn out: you get the iterator once, then keep looping back for another next until the first StopIteration breaks you out.

flowchart TD
    A["it = iter(numbers)"] --> B["call next(it)"]
    B --> C{"StopIteration?"}
    C -- "no: got a value" --> D["run loop body"]
    D --> B
    C -- "yes" --> E["loop ends"]

enumerate and zip

These two built-ins are iterables that make common loops cleaner, and they work by the same protocol above.

enumerate pairs each item with its position number (its index, counting from 0). Use it when you need both the item and where it sits:

colors = ["red", "green", "blue"]
for i, color in enumerate(colors):
    print(i, color)
# -> 0 red
# -> 1 green
# -> 2 blue

Here i, color unpacks each pair into two names at once. You can start counting from a different number with enumerate(colors, start=1).

zip walks several iterables in step, pairing up their items by position:

names = ["Ann", "Bo", "Cy"]
ages = [30, 25, 41]
for name, age in zip(names, ages):
    print(name, age)
# -> Ann 30
# -> Bo 25
# -> Cy 41

zip stops as soon as the shortest input runs out, so if ages had only two values, only two pairs would be produced and the extra name would be dropped. Both enumerate and zip produce their pairs lazily (one at a time, on demand), so they use O(1) extra space on top of the inputs; a full loop over them is O(n) time.

They combine naturally with comprehensions:

names = ["Ann", "Bo", "Cy"]
ages = [30, 25, 41]
people = {name: age for name, age in zip(names, ages)}
print(people)  # -> {'Ann': 30, 'Bo': 25, 'Cy': 41}

Common pitfalls

An iterator is one-shot. Once an iterator reaches the end, it stays exhausted. Asking again yields nothing:

it = iter([1, 2])
print(list(it))  # -> [1, 2]
print(list(it))  # -> []   (already used up)

list(it) drains the iterator by pulling items until StopIteration. The second call finds it empty. If you need to loop twice, iterate over the original list (which can produce a fresh iterator each time), not over a saved iterator.

zip silently truncates. If your inputs have different lengths, zip uses the shortest and drops the rest with no warning. Check lengths first if that would be a bug.

Do not modify a list while looping over it. Adding or removing items from a collection during its own for loop can skip elements or raise an error, because the iterator’s position no longer lines up with the changed contents. Build a new list with a comprehension instead:

numbers = [1, 2, 3, 4]
kept = [n for n in numbers if n != 2]  # safe: builds a new list
print(kept)  # -> [1, 3, 4]

{} is a dict, not a set. For an empty set, write set().

Practice

  1. Given words = ["apple", "kiwi", "banana", "fig"], use a list comprehension to build a list of the words that have more than four letters.

  2. Given the same words, use a dict comprehension to build a dictionary mapping each word to its length, but only for words that start with a letter before “m” in the alphabet (hint: compare word[0] < "m").

  3. Without using a for loop, use iter and next to print the first two items of [100, 200, 300, 400], then explain in a comment what a third next call would return and why.

Report a bug