Backend / Python core / Stdlib / 02_itertools.md

itertools — efficient iterator combinators

Updated 3 min read source
On this page11
  1. Infinite iterators
  2. Slicing iterators
  3. Combining iterables
  4. Filtering and selection
  5. Aggregation
  6. Grouping
  7. Combinatorics
  8. Splitting iterators with tee
  9. pairwise (3.10+)
  10. Common patterns
  11. Interview angle

itertools — efficient iterator combinators

itertools provides building blocks for working with iterators lazily and memory-efficiently. Most operations are O(1) memory regardless of input size.

Infinite iterators

python 
import itertools as it

it.count(start=0, step=1)        # 0, 1, 2, 3, ...
it.cycle([1, 2, 3])              # 1, 2, 3, 1, 2, 3, ...
it.repeat("hi", times=3)         # 'hi', 'hi', 'hi'
it.repeat(0)                     # 0, 0, 0, ...    (no `times` = infinite)

Use with islice to bound them. it.count() paired with zip is the modern enumerate analog.

Slicing iterators

python 
list(it.islice(it.count(), 5))       # [0, 1, 2, 3, 4]
list(it.islice(it.count(), 5, 10))   # [5, 6, 7, 8, 9]

Like list slicing but works on any iterable, including infinite ones.

Combining iterables

python 
list(it.chain([1, 2], [3, 4], [5]))         # [1, 2, 3, 4, 5]
list(it.chain.from_iterable([[1, 2], [3, 4]]))   # [1, 2, 3, 4]  ← flattens

list(zip([1, 2, 3], ['a', 'b', 'c']))       # [(1, 'a'), (2, 'b'), (3, 'c')]
list(it.zip_longest([1, 2, 3], ['a'], fillvalue='?'))   # [(1, 'a'), (2, '?'), (3, '?')]

chain.from_iterable is the fast way to flatten one level of nesting.

Filtering and selection

python 
list(it.compress("ABCDEF", [1, 0, 1, 0, 1, 1]))   # ['A', 'C', 'E', 'F']
list(it.takewhile(lambda x: x < 5, [1, 4, 6, 4, 1]))   # [1, 4]
list(it.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1]))   # [6, 4, 1]
list(it.filterfalse(lambda x: x % 2, range(10)))       # [0, 2, 4, 6, 8]

Aggregation

python 
list(it.accumulate([1, 2, 3, 4]))                  # [1, 3, 6, 10] — running sum
list(it.accumulate([1, 2, 3, 4], operator.mul))    # [1, 2, 6, 24] — running product
list(it.accumulate([5, 3, 7, 2], max))             # [5, 5, 7, 7] — running max

accumulate is great for prefix sums (DP problems, time series).

Grouping

python 
data = [("a", 1), ("a", 2), ("b", 3), ("b", 4), ("a", 5)]

for key, group in it.groupby(data, key=lambda x: x[0]):
    print(key, list(group))
# a [('a', 1), ('a', 2)]
# b [('b', 3), ('b', 4)]
# a [('a', 5)]            ← gotcha: groupby groups CONSECUTIVE elements only

Gotcha: groupby only groups adjacent equal keys. To group all matches by key, sort first:

python 
data.sort(key=lambda x: x[0])
for key, group in it.groupby(data, key=lambda x: x[0]):
    ...

Combinatorics

python 
list(it.product([1, 2], ['a', 'b']))            # [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
list(it.product(range(2), repeat=3))             # all 3-bit binary tuples (0,0,0)...(1,1,1)

list(it.permutations([1, 2, 3]))                 # 6 orderings
list(it.permutations([1, 2, 3], 2))              # 6 ordered pairs

list(it.combinations([1, 2, 3, 4], 2))           # [(1,2), (1,3), (1,4), (2,3), (2,4), (3,4)]
list(it.combinations_with_replacement([1, 2], 2))  # [(1,1), (1,2), (2,2)]

These materialize into iterators — be careful with large inputs (permutations(range(20)) is 20! items).

Splitting iterators with tee

python 
a, b = it.tee([1, 2, 3, 4], 2)
sum(a) + max(b)         # 10 + 4 = 14

tee buffers as much as needed for the slowest consumer. Don’t use it if one branch will lag far behind — memory grows.

pairwise (3.10+)

python 
# [(1, 2), (2, 3), (3, 4)]
list(it.pairwise([1, 2, 3, 4]))

Useful for diffing consecutive elements (sliding window of size 2).

Common patterns

python 
# Flatten one level
flat = list(it.chain.from_iterable(nested))

# Sliding window of size n (3.10+):
def windows(iterable, n):
    iters = it.tee(iterable, n)
    for i, it_ in enumerate(iters):
        next(it.islice(it_, i, i), None)
    return zip(*iters)

# [(1,2,3), (2,3,4), (3,4,5)]
list(windows([1, 2, 3, 4, 5], 3))

# Chunk into groups of n:
def chunked(iterable, n):
    it_ = iter(iterable)
    while chunk := list(it.islice(it_, n)):
        yield chunk

# [[0,1,2], [3,4,5], [6,7,8], [9]]
list(chunked(range(10), 3))

In Python 3.12+, itertools.batched does the chunking for you.

Interview angle

“Find all unique pairs in a list” → combinations. “Compute running sums” → accumulate. “Group consecutive duplicates” → groupby. “Why is chain.from_iterable faster than [x for sub in nested for x in sub]?” (Pure C; no Python-level loop.)

Previous
collections — specialized containers