Skip to content
intermediate Phase 4 · Python Functional Programming

itertools & functools

Leverage itertools for infinite iterators and functools for higher-order functions.

1h
0 problems
Topic Progress 0%

Infinite Iterators

count()

from itertools import count

# count(start=0, step=1)
counter = count(10, 2)  # 10, 12, 14, 16, ...

# Take first 5 values
from itertools import islice
first_5 = list(islice(counter, 5))
print(first_5)  # [10, 12, 14, 16, 18]

# Use with zip
numbers = [1, 2, 3, 4, 5]
numbered = list(zip(count(1), numbers))
print(numbered)  # [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]

cycle()

from itertools import cycle, islice

# cycle(iterable)
colors = cycle(['red', 'green', 'blue'])

# Take first 6 values
first_6 = list(islice(colors, 6))
print(first_6)  # ['red', 'green', 'blue', 'red', 'green', 'blue']

# Use for round-robin
def round_robin(*iterables):
    pending = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while pending:
        try:
            for next_func in nexts:
                yield next_func()
        except StopIteration:
            pending -= 1
            nexts = cycle(islice(nexts, pending))

repeat()

from itertools import repeat

# repeat(object, times=None)
# Finite repeat
ones = list(repeat(1, 5))
print(ones)  # [1, 1, 1, 1, 1]

# Infinite repeat (use with islice)
inf_ones = repeat(1)  # Infinite 1s
first_5 = list(islice(inf_ones, 5))
print(first_5)  # [1, 1, 1, 1, 1]

# Use with map
result = list(map(pow, range(5), repeat(2)))
print(result)  # [0, 1, 4, 9, 16]
# Equivalent to: [x**2 for x in range(5)]

Stopping Infinite Iterators

from itertools import count, islice, takewhile

# Method 1: islice
counter = count(1)
first_10 = list(islice(counter, 10))

# Method 2: takewhile
counter = count(1)
under_10 = list(takewhile(lambda x: x < 10, counter))
print(under_10)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

Combinatoric Iterators

permutations()

from itertools import permutations

# permutations(iterable, r=None)
items = ['A', 'B', 'C']

# All permutations of length 3
all_perms = list(permutations(items))
print(all_perms)  # [('A','B','C'), ('A','C','B'), ('B','A','C'), ...]
print(len(all_perms))  # 6 (3!)

# Permutations of length 2
perms_2 = list(permutations(items, 2))
print(perms_2)  # [('A','B'), ('A','C'), ('B','A'), ('B','C'), ('C','A'), ('C','B')]
print(len(perms_2))  # 6 (3P2)

combinations()

from itertools import combinations

# combinations(iterable, r)
items = ['A', 'B', 'C', 'D']

# Combinations of length 2
combs_2 = list(combinations(items, 2))
print(combs_2)  # [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')]
print(len(combs_2))  # 6 (4C2)

# Combinations of length 3
combs_3 = list(combinations(items, 3))nprint(combs_3)  # [('A','B','C'), ('A','B','D'), ('A','C','D'), ('B','C','D')]

combinations_with_replacement()

from itertools import combinations_with_replacement

# combinations_with_replacement(iterable, r)
items = ['A', 'B', 'C']

# With replacement
cwr = list(combinations_with_replacement(items, 2))
print(cwr)  # [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')]
print(len(cwr))  # 6 (n+r-1)Cr)

Product (Cartesian Product)

from itertools import product

# product(*iterables, repeat=1)
colors = ['red', 'blue']
sizes = ['S', 'M', 'L']

# Cartesian product
products = list(product(colors, sizes))
print(products)  # [('red','S'), ('red','M'), ('red','L'), ('blue','S'), ...]

# With repeat (power/permutations with repetition)
coins = [0, 1]  # 0=Heads, 1=Tails
outcomes = list(product(coins, repeat=3))
print(len(outcomes))  # 8 (2^3)

Combinatoric Formulas

# Permutations: P(n, r) = n! / (n-r)!
# Combinations: C(n, r) = n! / (r! * (n-r)!)

import math

n, r = 5, 3
perms = math.perm(n, r)  # 60
combs = math.comb(n, r)  # 10

print(f'P({n},{r}) = {perms}')  # P(5,3) = 60
print(f'C({n},{r}) = {combs}')  # C(5,3) = 10

Grouping & Chaining

groupby()

from itertools import groupby

# groupby(iterable, key=None)
# IMPORTANT: Data must be sorted by key!

# Sort by key first
items = [
    {'type': 'fruit', 'name': 'apple'},
    {'type': 'fruit', 'name': 'banana'},
    {'type': 'vegetable', 'name': 'carrot'},
    {'type': 'vegetable', 'name': 'daikon'}
]

# Sort by type (required for groupby!)
items.sort(key=lambda x: x['type'])

# Group by type
for key, group in groupby(items, key=lambda x: x['type']):
    print(f'{key}: {[item["name"] for item in group]}')
# fruit: ['apple', 'banana']
# vegetable: ['carrot', 'daikon']

# Simple example
words = ['apple', 'avocado', 'banana', 'blueberry', 'cherry']
words.sort()  # Sort by first letter

for letter, group in groupby(words, key=lambda w: w[0]):
    print(f'{letter}: {list(group)}')
# a: ['apple', 'avocado']
# b: ['banana', 'blueberry']
# c: ['cherry']

chain()

from itertools import chain

# chain(*iterables)
list1 = [1, 2, 3]
list2 = [4, 5, 6]
list3 = [7, 8, 9]

# Concatenate iterables
result = list(chain(list1, list2, list3))
print(result)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

# chain.from_iterable
lists = [[1, 2], [3, 4], [5, 6]]
result = list(chain.from_iterable(lists))
print(result)  # [1, 2, 3, 4, 5, 6]

# Equivalent to nested comprehension
result = [x for sublist in lists for x in sublist]

accumulate()

from itertools import accumulate

# accumulate(iterable, func=operator.add)
numbers = [1, 2, 3, 4, 5]

# Running sum (default)
result = list(accumulate(numbers))
print(result)  # [1, 3, 6, 10, 15]

# Running product
import operator
result = list(accumulate(numbers, operator.mul))
print(result)  # [1, 2, 6, 24, 120]

# Running max
result = list(accumulate(numbers, max))
print(result)  # [1, 2, 3, 4, 5]

# With initial value
result = list(accumulate(numbers, operator.add, initial=100))
print(result)  # [100, 101, 103, 106, 110, 115]

Practical Examples

from itertools import chain, groupby, accumulate

# Flatten and sort
nested = [[3, 1, 2], [6, 4, 5]]
flat_sorted = sorted(chain.from_iterable(nested))
print(flat_sorted)  # [1, 2, 3, 4, 5, 6]

# Chunked processing
def chunked(iterable, n):
    from itertools import islice
    iterator = iter(iterable)
    while True:
        chunk = list(islice(iterator, n))
        if not chunk:
            break
        yield chunk

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