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]]