Iterators and Comprehensions in Python

Posted on

Python has concise, expressive ways to loop and build collections. Comprehensions and the built-in iteration helpers make algorithm code shorter and easier to read, which matters when you are focused on the logic.

List comprehensions#

A comprehension builds a list in one readable line. Compare the loop and the comprehension:

1# Loop version
2squares = []
3for n in range(5):
4    squares.append(n * n)
5
6# Comprehension version
7squares = [n * n for n in range(5)]
8print(squares)   # [0, 1, 4, 9, 16]

Add a condition to filter:

1evens = [n for n in range(10) if n % 2 == 0]
2print(evens)   # [0, 2, 4, 6, 8]

Dict and set comprehensions#

The same syntax builds dictionaries and sets:

1squared_map = {n: n * n for n in range(4)}
2print(squared_map)   # {0: 0, 1: 1, 2: 4, 3: 9}
3
4unique_lengths = {len(word) for word in ["a", "bb", "cc"]}
5print(unique_lengths)   # {1, 2}

enumerate: index and value together#

When you need both the position and the item, use enumerate instead of manually tracking an index:

1for i, letter in enumerate(["a", "b", "c"]):
2    print(i, letter)   # 0 a, 1 b, 2 c

This is cleaner than for i in range(len(...)) and appears in many array algorithms.

zip: iterate over two sequences at once#

1names = ["ana", "ben"]
2scores = [90, 85]
3for name, score in zip(names, scores):
4    print(name, score)   # ana 90, ben 85
5
6paired = dict(zip(names, scores))
7print(paired)   # {'ana': 90, 'ben': 85}

Iterators and lazy evaluation#

An iterable is anything you can loop over. Some helpers are lazy: they produce values one at a time instead of building a whole list, which saves memory on large inputs:

1gen = (n * n for n in range(1000000))  # generator, nothing computed yet
2print(next(gen))   # 0
3print(next(gen))   # 1

range itself is lazy, which is why range(1000000) does not build a giant list.

Handy built-ins for DSA#

1nums = [5, 2, 9, 1]
2print(sum(nums))          # 17
3print(min(nums), max(nums))  # 1 9
4print(sorted(nums))       # [1, 2, 5, 9]
5print(any(n > 8 for n in nums))   # True
6print(all(n > 0 for n in nums))   # True

Combining any/all with a generator lets you express “does any item satisfy this” in one line, without writing a loop and a flag.

Wrapping up the primer#

You now have the core Python for data structures and algorithms: sequences, hashing, strings, functions, recursion, classes, and Pythonic iteration. Head to the main algorithm lessons, keep the Big-O cheat sheet open, and start Day 1 of the curriculum.