Python itertools: Iterator Building Blocks

from itertools import count, cycle, repeat
 
# Count forever
for i in count(10, 2):  # 10, 12, 14, 16, ...
    if i > 20:
        break
    print(i)
 
# Cycle through items
colors = cycle(["red", "green", "blue"])
for _ in range(5):
    print(next(colors))  # red, green, blue, red, green
 
# Repeat value
for x in repeat("hello", 3):
    print(x)  # hello hello hello

from itertools import chain
 
# Chain iterables together
letters = chain("abc", "def", "ghi")
print(list(letters))  # ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
 
# Flatten nested iterables
nested = [[1, 2], [3, 4], [5, 6]]
flat = list(chain.from_iterable(nested))
print(flat)  # [1, 2, 3, 4, 5, 6]

from itertools import islice
 
# Take first N items
nums = range(100)
first_five = list(islice(nums, 5))  # [0, 1, 2, 3, 4]
 
# Skip and take
middle = list(islice(nums, 10, 15))  # [10, 11, 12, 13, 14]
 
# With step
every_third = list(islice(nums, 0, 10, 3))  # [0, 3, 6, 9]

from itertools import filterfalse, takewhile, dropwhile
 
nums = [1, 4, 6, 4, 1, 0, 3, 5]
 
# Filter out truthy values
zeros = list(filterfalse(bool, [1, 0, 2, 0, 3]))  # [0, 0]
 
# Take while condition is true
ascending = list(takewhile(lambda x: x < 5, nums))  # [1, 4]
 
# Drop while condition is true
after_small = list(dropwhile(lambda x: x < 5, nums))  # [6, 4, 1, 0, 3, 5]

from itertools import groupby
 
# Group consecutive items
data = "AAAABBBCCDA"
groups = [(k, list(g)) for k, g in groupby(data)]
# [('A', ['A', 'A', 'A', 'A']), ('B', ['B', 'B', 'B']), ...]
 
# Group by key function (data must be sorted by key first!)
people = [
    {"name": "Alice", "dept": "Engineering"},
    {"name": "Bob", "dept": "Engineering"},
    {"name": "Carol", "dept": "Sales"},
]
people.sort(key=lambda x: x["dept"])
 
for dept, members in groupby(people, key=lambda x: x["dept"]):
    print(f"{dept}: {[m['name'] for m in members]}")

from itertools import combinations, permutations, product
 
items = ["a", "b", "c"]
 
# Combinations (order doesn't matter)
print(list(combinations(items, 2)))
# [('a', 'b'), ('a', 'c'), ('b', 'c')]
 
# Permutations (order matters)
print(list(permutations(items, 2)))
# [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]
 
# Cartesian product
print(list(product([1, 2], ["a", "b"])))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
 
# Repeat for power
print(list(product([0, 1], repeat=3)))
# All 3-bit binary numbers

from itertools import combinations_with_replacement
 
items = ["a", "b", "c"]
print(list(combinations_with_replacement(items, 2)))
# [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]

from itertools import accumulate
import operator
 
nums = [1, 2, 3, 4, 5]
 
# Running sum
print(list(accumulate(nums)))  # [1, 3, 6, 10, 15]
 
# Running product
print(list(accumulate(nums, operator.mul)))  # [1, 2, 6, 24, 120]
 
# Running max
print(list(accumulate(nums, max)))  # [1, 2, 3, 4, 5]
 
# Custom function
print(list(accumulate(nums, lambda a, b: a + b * 2)))

from itertools import zip_longest
 
a = [1, 2, 3]
b = [4, 5]
 
# Regular zip stops at shortest
print(list(zip(a, b)))  # [(1, 4), (2, 5)]
 
# zip_longest fills missing values
print(list(zip_longest(a, b, fillvalue=0)))  # [(1, 4), (2, 5), (3, 0)]

from itertools import pairwise
 
items = [1, 2, 3, 4, 5]
print(list(pairwise(items)))
# [(1, 2), (2, 3), (3, 4), (4, 5)]
 
# Before 3.10:
def pairwise_compat(iterable):
    a, b = iter(iterable), iter(iterable)
    next(b, None)
    return zip(a, b)

from itertools import batched
 
items = range(10)
print(list(batched(items, 3)))
# [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9,)]
 
# Before 3.12:
def batched_compat(iterable, n):
    from itertools import islice
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch

from itertools import starmap
 
# Apply function to unpacked arguments
pairs = [(2, 3), (4, 5), (6, 7)]
print(list(starmap(pow, pairs)))  # [8, 1024, 279936]
 
# Same as:
# [pow(a, b) for a, b in pairs]

from itertools import tee
 
nums = iter([1, 2, 3, 4, 5])
a, b = tee(nums, 2)
 
print(list(a))  # [1, 2, 3, 4, 5]
print(list(b))  # [1, 2, 3, 4, 5]
 
# Warning: tee stores elements in memory
# Don't use original iterator after tee

from itertools import islice, chain, groupby, accumulate
 
# Sliding window
def sliding_window(iterable, n):
    from collections import deque
    it = iter(iterable)
    window = deque(islice(it, n), maxlen=n)
    if len(window) == n:
        yield tuple(window)
    for x in it:
        window.append(x)
        yield tuple(window)
 
# Chunked iteration
def chunked(iterable, n):
    it = iter(iterable)
    while chunk := tuple(islice(it, n)):
        yield chunk
 
# Flatten one level
def flatten(list_of_lists):
    return chain.from_iterable(list_of_lists)
 
# Run-length encoding
def run_length_encode(iterable):
    return [(k, len(list(g))) for k, g in groupby(iterable)]
 
# Unique elements preserving order
def unique(iterable):
    seen = set()
    for item in iterable:
        if item not in seen:
            seen.add(item)
            yield item
 
# First n items or less
def take(n, iterable):
    return list(islice(iterable, n))
 
# Nth item
def nth(iterable, n, default=None):
    return next(islice(iterable, n, None), default)

from itertools import chain
 
# BAD: Creates intermediate list
data = []
for chunk in chunks:
    data.extend(process(chunk))
 
# GOOD: Lazy iteration
data = chain.from_iterable(process(chunk) for chunk in chunks)
 
# BAD: Loads all into memory
sum([x ** 2 for x in range(10_000_000)])
 
# GOOD: Generator expression
sum(x ** 2 for x in range(10_000_000))

from itertools import count, takewhile
 
# Generate IDs
id_generator = count(1)
next_id = lambda: next(id_generator)
 
# Find first match
def first_where(iterable, predicate):
    return next((x for x in iterable if predicate(x)), None)
 
# Interleave sequences
def interleave(*iterables):
    from itertools import chain, zip_longest
    sentinel = object()
    for items in zip_longest(*iterables, fillvalue=sentinel):
        yield from (x for x in items if x is not sentinel)
 
# Partition into two groups
def partition(predicate, iterable):
    from itertools import filterfalse, tee
    t1, t2 = tee(iterable)
    return filterfalse(predicate, t1), filter(predicate, t2)