Módulo itertools do Python

O módulo itertools do Python é um kit de ferramentas da biblioteca padrão com blocos de construção rápidos e eficientes em memória para trabalhar com iteradores. Cada função em itertools retorna um iterador — ela produz valores sob demanda em vez de construir uma lista na memória — tornando o módulo ideal para grandes conjuntos de dados, sequências infinitas e pipelines de dados combináveis.

Este capítulo abrange todas as três categorias de funções do itertools: iteradores infinitos (count, cycle, repeat), iteradores combinatórios (product, permutations, combinations, combinations_with_replacement) e iteradores de terminação (chain, islice, groupby, compress, filterfalse, takewhile, dropwhile, starmap, zip_longest, accumulate, pairwise).

Nenhuma instalação é necessária — itertools é distribuído com toda instalação do Python 3:

import itertools

Por que usar itertools?

Considere ler os primeiros 10 múltiplos de um número. Sem itertools você precisa de uma lista ou um contador manual. Com itertools.count e itertools.islice a intenção fica imediatamente clara e o uso de memória permanece constante:

import itertools

multiples = itertools.islice(itertools.count(0, 7), 10)
print(list(multiples))
# [0, 7, 14, 21, 28, 35, 42, 49, 56, 63]

A filosofia do itertools: construa uma peça pequena e correta, depois combine-a com outras. Encadear duas funções do itertools é mais rápido e menos propenso a erros do que escrever o loop equivalente manualmente.

Iteradores Infinitos

Esses iteradores produzem valores indefinidamente. Sempre os combine com islice, um for … break, ou outro mecanismo de terminação para evitar um loop infinito.

count(start=0, step=1)

count produz uma sequência de números com espaçamento uniforme. É essencialmente um range sem limite superior e com suporte a floats e passos negativos.

import itertools

# Integer counter
for n in itertools.islice(itertools.count(10), 5):
    print(n, end=' ')
# 10 11 12 13 14

print()

# Float step
for n in itertools.islice(itertools.count(0.0, 0.5), 5):
    print(n, end=' ')
# 0.0 0.5 1.0 1.5 2.0

print()

# Countdown
for n in itertools.islice(itertools.count(100, -10), 5):
    print(n, end=' ')
# 100 90 80 70 60

count é útil quando você precisa numerar itens de um iterável sem saber quantos existem de antemão — o idioma do enumerate mas com início e passo personalizados.

cycle(iterable)

cycle repete os elementos de qualquer iterável indefinidamente.

import itertools

colours = itertools.cycle(['red', 'green', 'blue'])
for i, colour in enumerate(colours):
    if i == 7:
        break
    print(colour, end=' ')
# red green blue red green blue red

Uso prático — atribuindo itens a equipes em modo round-robin:

import itertools

teams = itertools.cycle(['Alpha', 'Beta', 'Gamma'])
players = ['Alice', 'Bob', 'Carol', 'Dave', 'Eve']
assignments = {player: team for player, team in zip(players, teams)}
print(assignments)
# {'Alice': 'Alpha', 'Bob': 'Beta', 'Carol': 'Gamma', 'Dave': 'Alpha', 'Eve': 'Beta'}

repeat(object, times=None)

repeat produz o mesmo object times vezes (ou para sempre se times for omitido).

import itertools

# Finite repeat
print(list(itertools.repeat('hello', 3)))
# ['hello', 'hello', 'hello']

# Used as a fixed argument supplier in map()
squares = list(map(pow, range(1, 6), itertools.repeat(2)))
print(squares)
# [1, 4, 9, 16, 25]

O padrão map(pow, range(1, 6), repeat(2)) é um idioma comum para fornecer um segundo argumento constante a uma função de dois argumentos.

Iteradores Combinatórios

Esses iteradores produzem todas as combinações, permutações ou produtos cartesianos de um iterável de entrada. São essenciais para buscas por força bruta, geração de casos de teste e problemas de combinatória.

product(*iterables, repeat=1)

product calcula o produto cartesiano — todas as combinações ordenadas onde um elemento é retirado de cada iterável. É equivalente a loops for aninhados.

import itertools

suits = ['Hearts', 'Diamonds']
ranks = ['A', 'K', 'Q']
deck = list(itertools.product(suits, ranks))
print(deck)
# [('Hearts', 'A'), ('Hearts', 'K'), ('Hearts', 'Q'),
#  ('Diamonds', 'A'), ('Diamonds', 'K'), ('Diamonds', 'Q')]

Use repeat para calcular o produto de um iterável consigo mesmo múltiplas vezes:

import itertools

# All 2-digit binary numbers
binary_pairs = list(itertools.product([0, 1], repeat=2))
print(binary_pairs)
# [(0, 0), (0, 1), (1, 0), (1, 1)]

Atenção: product materializa os iteráveis de entrada na memória (para permitir múltiplas passagens), portanto não passe iteradores enormes como entrada.

permutations(iterable, r=None)

permutations produz todos os arranjos ordenados de r elementos retirados da entrada. Quando r é omitido, todos os elementos são usados.

import itertools

# All orderings of 3 letters
perms = list(itertools.permutations('ABC'))
print(perms)
# [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'),
#  ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]
print(len(perms))   # 6  (3! = 6)

# 2-element permutations
perms2 = list(itertools.permutations('ABC', 2))
print(perms2)
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
print(len(perms2))  # 6  (3 * 2 = 6)

A ordem importa nas permutações — ('A', 'B') e ('B', 'A') são resultados distintos.

combinations(iterable, r)

combinations produz todas as seleções não ordenadas de r elementos. Ao contrário de permutations, a ordem não importa — cada subconjunto aparece apenas uma vez.

import itertools

# All 2-element subsets of [1, 2, 3, 4]
combos = list(itertools.combinations([1, 2, 3, 4], 2))
print(combos)
# [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
print(len(combos))  # 6  (C(4,2) = 6)

Um caso de uso comum — verificar todos os pares de itens quanto a uma propriedade:

import itertools

words = ['bat', 'tab', 'cat', 'tac']
anagram_pairs = [
    (a, b) for a, b in itertools.combinations(words, 2)
    if sorted(a) == sorted(b)
]
print(anagram_pairs)
# [('bat', 'tab'), ('cat', 'tac')]

combinations_with_replacement(iterable, r)

Como combinations, mas permite que cada elemento apareça mais de uma vez em uma seleção.

import itertools

# All 2-element combinations with repetition from [1, 2, 3]
combos = list(itertools.combinations_with_replacement([1, 2, 3], 2))
print(combos)
# [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

Isso é útil para gerar todos os resultados possíveis de dados, sequências de lançamento de moeda ou escolhas de caracteres para senhas.

Funções combinatórias em resumo

Iteradores de Terminação

Iteradores de terminação processam uma entrada finita e param quando essa entrada é esgotada.

chain(*iterables)

chain trata vários iteráveis como uma única sequência contínua sem construir uma nova lista.

import itertools

a = [1, 2, 3]
b = (4, 5)
c = range(6, 9)
combined = list(itertools.chain(a, b, c))
print(combined)
# [1, 2, 3, 4, 5, 6, 7, 8]

chain.from_iterable aceita um único iterável de iteráveis — útil quando você não sabe o número de sequências com antecedência:

import itertools

nested = [[1, 2], [3, 4], [5, 6]]
flat = list(itertools.chain.from_iterable(nested))
print(flat)
# [1, 2, 3, 4, 5, 6]

Esta é uma alternativa rápida e eficiente em memória para [item for sublist in nested for item in sublist].

islice(iterable, stop) / islice(iterable, start, stop, step=1)

islice fatia qualquer iterador — incluindo os infinitos — sem materializá-lo. Os argumentos espelham a notação de slice do Python, mas aceitam apenas inteiros não negativos.

import itertools

# First 5 elements
print(list(itertools.islice(range(100), 5)))
# [0, 1, 2, 3, 4]

# Elements 10–14 (start inclusive, stop exclusive)
print(list(itertools.islice(range(100), 10, 15)))
# [10, 11, 12, 13, 14]

# Every other element from position 0 to 10
print(list(itertools.islice(range(20), 0, 10, 2)))
# [0, 2, 4, 6, 8]

islice não suporta índices negativos nem passos negativos (ao contrário do fatiamento regular de listas).

groupby(iterable, key=None)

groupby agrupa elementos consecutivos que compartilham o mesmo valor de chave. Retorna pares (key, group_iterator).

import itertools

data = [
    ('fruit', 'apple'),
    ('fruit', 'banana'),
    ('veggie', 'carrot'),
    ('veggie', 'broccoli'),
    ('fruit', 'cherry'),
]

for category, group in itertools.groupby(data, key=lambda x: x[0]):
    items = [item[1] for item in group]
    print(f'{category}: {items}')
# fruit: ['apple', 'banana']
# veggie: ['carrot', 'broccoli']
# fruit: ['cherry']

Atenção importante: groupby só agrupa elementos consecutivos iguais. Se seus dados não estiverem pré-ordenados pela chave, itens semelhantes em posições diferentes formam grupos separados (como mostrado acima — 'cherry' inicia um novo grupo 'fruit' em vez de se juntar ao primeiro). Sempre ordene pela chave antes de chamar groupby:

import itertools

data = [
    ('fruit', 'apple'),
    ('veggie', 'carrot'),
    ('fruit', 'banana'),
    ('veggie', 'broccoli'),
    ('fruit', 'cherry'),
]

# Sort first, then group
sorted_data = sorted(data, key=lambda x: x[0])
for category, group in itertools.groupby(sorted_data, key=lambda x: x[0]):
    items = [item[1] for item in group]
    print(f'{category}: {items}')
# fruit: ['apple', 'banana', 'cherry']
# veggie: ['carrot', 'broccoli']

Note também que o iterador de grupo se torna inválido assim que você avança para a próxima chave — consuma cada grupo antes de chamar next() no iterador externo.

compress(data, selectors)

compress filtra data mantendo apenas os elementos cujo valor selector correspondente é verdadeiro.

import itertools

names = ['Alice', 'Bob', 'Carol', 'Dave', 'Eve']
active = [True, False, True, True, False]

result = list(itertools.compress(names, active))
print(result)
# ['Alice', 'Carol', 'Dave']

compress é equivalente a [d for d, s in zip(data, selectors) if s], mas é mais rápido e evita a lista intermediária.

filterfalse(predicate, iterable)

filterfalse é o complemento do filter embutido — ele produz elementos para os quais o predicado retorna False.

import itertools

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

# Keep only odd numbers (those that fail the even test)
odds = list(itertools.filterfalse(lambda x: x % 2 == 0, numbers))
print(odds)
# [1, 3, 5, 7, 9]

takewhile(predicate, iterable)

takewhile produz elementos enquanto o predicado for True, depois para imediatamente — mesmo que elementos posteriores satisfaçam o predicado.

import itertools

data = [2, 4, 6, 3, 8, 10]

# Stop as soon as an odd number appears
evens_from_start = list(itertools.takewhile(lambda x: x % 2 == 0, data))
print(evens_from_start)
# [2, 4, 6]

dropwhile(predicate, iterable)

dropwhile é o espelho de takewhile: ele pula elementos enquanto o predicado for True, depois produz todos os elementos restantes (incluindo aqueles onde o predicado seria True novamente).

import itertools

data = [2, 4, 6, 3, 8, 10]

# Drop leading even numbers, yield everything from the first odd onward
result = list(itertools.dropwhile(lambda x: x % 2 == 0, data))
print(result)
# [3, 8, 10]

takewhile e dropwhile são úteis para processar arquivos de log ou fluxos onde você quer pular uma seção de cabeçalho ou parar em uma linha sentinela.

starmap(function, iterable)

starmap aplica uma função a cada elemento de um iterável, desempacotando o elemento como argumentos posicionais. É o equivalente de map para iteráveis de tuplas.

import itertools

pairs = [(2, 3), (4, 2), (10, 3)]
results = list(itertools.starmap(pow, pairs))
print(results)
# [8, 16, 1000]

Compare com map(pow, [2, 4, 10], [3, 2, 3]) — starmap funciona quando seus argumentos já estão agrupados como tuplas.

zip_longest(*iterables, fillvalue=None)

O zip embutido para no iterável mais curto. zip_longest preenche iteráveis mais curtos com fillvalue para que todos os iteráveis sejam completamente consumidos.

import itertools

a = [1, 2, 3]
b = ['a', 'b', 'c', 'd', 'e']

print(list(zip(a, b)))
# [(1, 'a'), (2, 'b'), (3, 'c')]  — b's 'd' and 'e' are lost

print(list(itertools.zip_longest(a, b, fillvalue=0)))
# [(1, 'a'), (2, 'b'), (3, 'c'), (0, 'd'), (0, 'e')]

accumulate(iterable, func=operator.add, *, initial=None)

accumulate calcula totais acumulados (ou qualquer outra agregação acumulada). Por padrão soma, mas você pode passar qualquer função de dois argumentos.

import itertools
import operator

numbers = [1, 2, 3, 4, 5]

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

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

# Running maximum
data = [3, 1, 4, 1, 5, 9, 2, 6]
print(list(itertools.accumulate(data, max)))
# [3, 3, 4, 4, 5, 9, 9, 9]

O parâmetro initial (Python 3.8+) adiciona um valor inicial antes do primeiro elemento:

import itertools

print(list(itertools.accumulate([1, 2, 3], initial=100)))
# [100, 101, 103, 106]

pairwise(iterable)

pairwise (Python 3.10+) produz pares sobrepostos consecutivos do iterável.

import itertools

data = [1, 2, 3, 4, 5]
print(list(itertools.pairwise(data)))
# [(1, 2), (2, 3), (3, 4), (4, 5)]

Isso é útil para calcular diferenças entre valores consecutivos ou para lógica de janela deslizante onde o tamanho da janela é exatamente 2:

import itertools

prices = [10.0, 12.5, 11.0, 13.5, 15.0]
changes = [b - a for a, b in itertools.pairwise(prices)]
print(changes)
# [2.5, -1.5, 2.5, 1.5]

Antes do Python 3.10, o equivalente era zip(data, data[1:]) (funciona para sequências) ou uma abordagem manual baseada em tee (funciona para iteradores arbitrários).

Compondo itertools em Pipelines

O real poder do itertools emerge quando você combina funções. Como toda função retorna um iterador, você pode encadeá-las com zero listas intermediárias.

Exemplo: frequências das 3 palavras mais comuns em um texto

import itertools
import operator

text = "the quick brown fox jumps over the lazy dog the fox"
words = text.split()

# Sort words so groupby can collect identical words together
sorted_words = sorted(words)

# Count each word using groupby
word_counts = (
    (key, sum(1 for _ in group))
    for key, group in itertools.groupby(sorted_words)
)

# Sort by count descending, take the top 3
top3 = list(itertools.islice(
    sorted(word_counts, key=operator.itemgetter(1), reverse=True),
    3
))
print(top3)
# [('the', 3), ('fox', 2), ('brown', 1)]

Exemplo: dividindo um iterável em blocos de tamanho fixo

import itertools

def batched(iterable, n):
    """Yield successive n-sized tuples from iterable."""
    it = iter(iterable)
    while chunk := tuple(itertools.islice(it, n)):
        yield chunk

data = range(10)
for batch in batched(data, 3):
    print(batch)
# (0, 1, 2)
# (3, 4, 5)
# (6, 7, 8)
# (9,)

O Python 3.12 inclui itertools.batched como função embutida, então você pode substituir o auxiliar acima por itertools.batched(data, 3) no Python moderno.

Referência Rápida

Para os conceitos de avaliação preguiçosa por trás do itertools, consulte Python Generators e Python Iterators. Para auxiliares de estilo funcional que complementam o itertools, veja Python Lambda Functions e o Python collections Module.

Prática