C.W.K.
Stream
Lesson 03 of 07 · published

itertools — The Lazy Iteration Toolbox

~25 min · itertools, chain, groupby, tee, product

Level 0Curious
0 XP0/93 lessons0/23 achievements
0/100 XP to next level100 XP to go0% complete

The standard library's iteration superpowers

Once you understand iterators and generators, itertools becomes one of the most useful modules in the standard library. It's a collection of fast, memory-efficient building blocks for working with iterables. Unlike a comprehension that builds a list, an itertools function returns a lazy iterator. You compose them like a pipeline.

The infinite trio — count, cycle, repeat

count(start, step) yields start, start+step, ... forever. cycle(iter) yields elements of iter over and over. repeat(value, times) yields the same value, optionally a fixed number of times. Always pair these with something that stops (islice, a break, zip against a finite iterable) — otherwise you have an infinite loop.

chain, zip_longest, islice — the structural ones

chain(a, b, c) yields all elements of a, then b, then c — like concatenation but lazy. zip_longest is like zip but doesn't stop at the shortest iterable; it pads with a fill value. islice(iter, start, stop, step) is slicing for any iterable, including infinite ones.

groupby — the one that surprises everyone

groupby(iter, key) groups consecutive elements that share a key. Two things to know: it only groups consecutive elements (sort first if you want all groups together), and the group it yields is itself an iterator that you must consume before moving on. People expect it to behave like SQL GROUP BY; it's closer to Unix uniq.

tee — splitting an iterator

An iterator can only be consumed once. tee(iter, n) creates n independent iterators from one source — internally it buffers as needed. Use sparingly: if one branch races far ahead of the others, the buffer grows. For most cases where you'd reach for tee, materializing the source as a list is simpler.

product, permutations, combinations

The combinatoric trio. product(a, b) is the Cartesian product — every (x, y) pair. permutations(iter, r) yields all ordered r-length arrangements. combinations(iter, r) yields unordered r-length selections. They replace nested for-loops for the entire genre of "all combinations of..." problems.

Code

Infinite iterators — pair with something that stops·python
import itertools as it

# count — like range but unbounded
for n in it.count(10, 2):
    if n > 20:
        break
    print(n)              # 10 12 14 16 18 20

# cycle — round-robin
colors = it.cycle(["red", "green", "blue"])
for _ in range(7):
    print(next(colors))
# red green blue red green blue red

# repeat — limited or infinite
print(list(it.repeat("x", 4)))    # ['x', 'x', 'x', 'x']

# Pair an infinite zip with a finite list
pairs = list(zip(it.count(1), ["a", "b", "c"]))
print(pairs)                       # [(1, 'a'), (2, 'b'), (3, 'c')]
chain, zip_longest, islice — structural·python
import itertools as it

# chain — concatenate iterables
result = list(it.chain([1, 2], [3, 4], [5]))
print(result)                      # [1, 2, 3, 4, 5]

# zip_longest — don't stop at the shortest
result = list(it.zip_longest([1, 2, 3], ["a", "b"], fillvalue="?"))
print(result)                      # [(1, 'a'), (2, 'b'), (3, '?')]

# islice — slice ANY iterable, including infinite
def naturals():
    n = 1
    while True:
        yield n
        n += 1

first_5 = list(it.islice(naturals(), 5))
print(first_5)                     # [1, 2, 3, 4, 5]
groupby — only groups CONSECUTIVE·python
import itertools as it

items = ["apple", "ant", "banana", "berry", "apricot"]

# Naive groupby — only groups consecutive matches
for key, group in it.groupby(items, key=lambda x: x[0]):
    print(key, list(group))
# a ['apple', 'ant']
# b ['banana', 'berry']
# a ['apricot']               <- second 'a' group, NOT merged

# Sort first for SQL-style grouping
items.sort(key=lambda x: x[0])
for key, group in it.groupby(items, key=lambda x: x[0]):
    print(key, list(group))
# a ['apple', 'ant', 'apricot']
# b ['banana', 'berry']
product / permutations / combinations·python
import itertools as it

# product — Cartesian product (replaces nested for-loops)
for x, y in it.product([1, 2], ["a", "b"]):
    print(x, y)
# 1 a
# 1 b
# 2 a
# 2 b

# permutations — all ORDERED arrangements
print(list(it.permutations("abc", 2)))
# [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]

# combinations — UNORDERED selections
print(list(it.combinations("abc", 2)))
# [('a', 'b'), ('a', 'c'), ('b', 'c')]

# combinations_with_replacement — selections that can reuse
print(list(it.combinations_with_replacement("ab", 2)))
# [('a', 'a'), ('a', 'b'), ('b', 'b')]
tee — splitting one iterator into many·python
import itertools as it

source = (x*x for x in range(5))
a, b, c = it.tee(source, 3)

print(list(a))               # [0, 1, 4, 9, 16]
print(list(b))               # [0, 1, 4, 9, 16]
print(list(c))               # [0, 1, 4, 9, 16]

# Note — DON'T use the original `source` after tee()
# It's fed into the tee buffers

# For most cases, just materializing as a list is clearer:
squares = [x*x for x in range(5)]

External links

Exercise

Given events = [{"user": "a", "action": "login"}, {"user": "a", "action": "click"}, {"user": "b", "action": "login"}, {"user": "b", "action": "logout"}, {"user": "a", "action": "logout"}], use itertools.groupby (with appropriate sorting) to produce a dict {user: [actions...]}. Then use itertools.product to generate all (user, color) pairs where colors is ["red", "green", "blue"]. Then use islice with an infinite count() to produce [10, 13, 16, 19, 22].

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.