`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

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

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

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

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

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

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:

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

Combinatorics

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`

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+)

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

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

Common patterns

# 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)

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

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

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

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.)

itertools — efficient iterator combinators