Python O Módulo itertools

Nas aulas Iteradores e Geradores, você viu que um iterador entrega um item por vez, sob encomenda. O módulo itertools ("ferramentas de iteração") é uma caixa de iteradores prontos para os trabalhos que aparecem sempre: numerar sem fim, revezar, emendar listas, agrupar, montar todas as combinações de um cardápio. Cada ferramenta substitui um laço que você escreveria na mão, e quase sempre com menos chance de errar.

A ideia: as máquinas da cozinha

A Ana poderia bater massa com um garfo, mas a padaria tem batedeira, liquidificador e máquina de fatiar. Cada máquina faz um serviço só, e faz bem. As ferramentas do itertools são assim: cada uma recebe iteráveis (listas, textos, geradores) e devolve um iterador que entrega o resultado aos poucos. Como iteradores, elas gastam pouca memória, e se gastam numa passada só.

Importe só o que for usar: from itertools import count, cycle.

1. count: contar sem fim

count(inicio) entrega inicio, inicio + 1, inicio + 2… para sempre. Sozinho, ele nunca para; junto com algo que tem fim, como o zip (aula enumerate() e zip()), ele para quando o outro acaba:

from itertools import count, islice

clientes = ["Ana", "Bruno", "Carla"]
for senha, nome in zip(count(101), clientes):
    print(f"Senha {senha}: {nome}")

print(list(islice(count(10, 5), 4)))
print(count(1))

$ python3 contar.py
Senha 101: Ana
Senha 102: Bruno
Senha 103: Carla
[10, 15, 20, 25]
count(1)

zip(count(101), clientes)

O zip anda nos dois ao mesmo tempo e para quando o mais curto acaba. Os clientes acabam em 3; o count nunca chega a dar o 104.

count(10, 5)

Começa em 10 e anda de 5 em 5. O segundo argumento é o passo.

islice(…, 4)

"Fatia de iterador": só os 4 primeiros. É o jeito de cortar um iterador sem fim (a aula Geradores já o usou).

print(count(1))

Mostra só count(1): o iterador não calculou nada ainda.

Para numerar uma lista que já existe, o enumerate(clientes, start=101) faz o mesmo. O count serve quando o que vai ser numerado não é uma lista pronta, ou quando você quer passos diferentes de 1.

2. cycle: em roda

cycle(sequencia) entrega os itens da sequência e, quando acaba, recomeça do primeiro, para sempre. O rodízio de quem abre a padaria:

from itertools import cycle

dias = ["seg", "ter", "qua", "qui", "sex"]
equipe = cycle(["Ana", "Bruno"])

for dia, quem in zip(dias, equipe):
    print(dia, "-", quem, "abre a padaria")

$ python3 ciclo.py
seg - Ana abre a padaria
ter - Bruno abre a padaria
qua - Ana abre a padaria
qui - Bruno abre a padaria
sex - Ana abre a padaria

Mais uma vez, quem dá o fim é o zip: os 5 dias. Sem o dias, um for quem in equipe: nunca terminaria.

3. chain: emendar sequências

chain(a, b, c) ("corrente") percorre a inteira, depois b, depois c, como se fossem uma lista só, mas sem criar a lista grande:

from itertools import chain

manha = ["pão", "leite"]
tarde = ["café", "bolo"]
noite = ["pizza"]

for item in chain(manha, tarde, noite):
    print(item, end=" ")
print()

pedidos_por_mesa = [["suco", "pastel"], ["café"], ["coxinha", "refrigerante"]]
print(list(chain.from_iterable(pedidos_por_mesa)))
print(manha + tarde + noite)

$ python3 corrente.py
pão leite café bolo pizza
['suco', 'pastel', 'café', 'coxinha', 'refrigerante']
['pão', 'leite', 'café', 'bolo', 'pizza']

• chain(manha, tarde, noite): as três listas em sequência. O end=" " põe tudo na mesma linha (aula As Opções sep e end).
• chain.from_iterable(lista_de_listas): para quando as listas já estão dentro de uma lista. "Achata" um nível.
• manha + tarde + noite dá o mesmo resultado, mas monta uma lista nova com tudo. O chain funciona também com geradores e arquivos, que não aceitam +.

4. groupby: juntar vizinhos iguais

A padaria anotou as vendas do dia, uma por linha, e quer o total de cada produto. O groupby ("agrupar por") junta itens seguidos que têm a mesma chave:

from itertools import groupby

vendas = [("leite", 3), ("pão", 12), ("café", 1), ("pão", 8), ("leite", 2)]

vendas.sort(key=lambda venda: venda[0])
print(vendas)

for produto, grupo in groupby(vendas, key=lambda venda: venda[0]):
    quantidades = [qtd for _, qtd in grupo]
    print(f"{produto}: {quantidades} -> {sum(quantidades)}")

$ python3 agrupar.py
[('café', 1), ('leite', 3), ('leite', 2), ('pão', 12), ('pão', 8)]
café: [1] -> 1
leite: [3, 2] -> 5
pão: [12, 8] -> 20

vendas.sort(key=…)

Primeiro, ordena pela mesma chave que vai agrupar (aula Ordenando Listas). Agora as vendas de leite estão lado a lado, e as de pão também.

key=lambda venda: venda[0]

A chave do grupo: o nome do produto, posição 0 da tupla (aula lambda).

for produto, grupo in groupby(…)

A cada volta, produto é a chave e grupo é um iterador com os itens daquele grupo.

[qtd for _, qtd in grupo]

Tira as quantidades do grupo. O _ é o nome do produto, que aqui não interessa.

A pegadinha: sem ordenar, os grupos se repetem

O groupby só olha o vizinho. Ele não procura a lista inteira: quando a chave muda, fecha o grupo e abre outro. Veja o mesmo programa sem o sort:

from itertools import groupby

vendas = [("leite", 3), ("pão", 12), ("café", 1), ("pão", 8), ("leite", 2)]

for produto, grupo in groupby(vendas, key=lambda venda: venda[0]):
    quantidades = [qtd for _, qtd in grupo]
    print(f"{produto}: {quantidades} -> {sum(quantidades)}")

$ python3 agrupar_errado.py
leite: [3] -> 3
pão: [12] -> 12
café: [1] -> 1
pão: [8] -> 8
leite: [2] -> 2

Nenhum erro, nenhum aviso, e a resposta está errada: o pão aparece duas vezes, com 12 e com 8, em vez de 20. Cada vez que o produto mudou de uma linha para a outra, começou um grupo novo. É como uma pessoa que separa as fichas olhando só a de cima da pilha: se as fichas estão embaralhadas, ela faz pilhas repetidas.

Esse comportamento é útil quando você quer justamente as "corridas" de itens iguais, como numa contagem de letras repetidas em sequência:

from itertools import groupby

print([(letra, len(list(grupo))) for letra, grupo in groupby("aaabbcaa")])

$ python3 agrupar_letras.py
[('a', 3), ('b', 2), ('c', 1), ('a', 2)]

Os dois a do fim formam um grupo separado, porque não estão encostados nos três primeiros. Para somar todos os itens de cada chave, onde quer que estejam, ordene antes; ou use um defaultdict(int) ou um Counter, da aula O Módulo collections, que não precisam de ordem.

from itertools import groupby

vendas = [("leite", 3), ("pão", 12), ("café", 1), ("pão", 8), ("leite", 2)]

for produto, grupo in groupby(vendas, key=lambda venda: venda[0]):
    total = sum(qtd for _, qtd in grupo)
    print(f"{produto}: {total}")

from itertools import groupby

vendas = [("leite", 3), ("pão", 12), ("café", 1), ("pão", 8), ("leite", 2)]
vendas.sort()

for produto, grupo in groupby(vendas, key=lambda venda: venda[0]):
    total = sum(qtd for _, qtd in grupo)
    print(f"{produto}: {total}")

5. Combinações: combinations, permutations e product

Três ferramentas montam "todas as possibilidades". A diferença entre elas é a pergunta: a ordem importa? e de onde vêm os itens?

from itertools import combinations, permutations

sabores = ["chocolate", "morango", "limão"]

print("Casquinha com 2 bolas (a ordem não importa):")
for dupla in combinations(sabores, 2):
    print(" ", dupla)

print("Pódio do concurso de bolos (a ordem importa):")
for primeiro, segundo in permutations(["Ana", "Bruno", "Carla"], 2):
    print(f"  1º {primeiro}, 2º {segundo}")

$ python3 combinar.py
Casquinha com 2 bolas (a ordem não importa):
  ('chocolate', 'morango')
  ('chocolate', 'limão')
  ('morango', 'limão')
Pódio do concurso de bolos (a ordem importa):
  1º Ana, 2º Bruno
  1º Ana, 2º Carla
  1º Bruno, 2º Ana
  1º Bruno, 2º Carla
  1º Carla, 2º Ana
  1º Carla, 2º Bruno

• combinations(sabores, 2): grupos de 2, sem repetir e sem ligar para a ordem. Chocolate com morango é a mesma casquinha que morango com chocolate, então só aparece uma vez: 3 casquinhas.
• permutations(pessoas, 2): arrumações de 2 em que a ordem importa. Ana em 1º e Bruno em 2º é outro pódio que Bruno em 1º e Ana em 2º: 6 pódios.

Quantas são? O math (aula O Módulo math) tem as contas prontas, comb e perm, que dão o número sem montar a lista:

from itertools import combinations, permutations
from math import comb, perm

turma = ["Ana", "Bruno", "Carla", "Davi", "Eva"]
print(len(list(combinations(turma, 2))), comb(5, 2))
print(len(list(permutations(turma, 2))), perm(5, 2))
print(len(list(permutations(turma))))

$ python3 contas.py
10 10
20 20
120

Sem o segundo argumento, permutations(turma) dá todas as filas possíveis com os 5: 120. Com 10 pessoas já seriam 3.628.800, e com 15, mais de um trilhão: essas listas crescem muito depressa, então pense antes de pedir list() delas.

product: um de cada

product(a, b) ("produto") combina um item de cada sequência, como dois for um dentro do outro (aula Laços Aninhados). O almoço do restaurante, prato + bebida:

from itertools import product

pratos = ["arroz e feijão", "macarrão"]
bebidas = ["suco", "refrigerante", "água"]

for prato, bebida in product(pratos, bebidas):
    print(f"{prato} + {bebida}")

print(list(product([0, 1], repeat=3)))

$ python3 produto.py
arroz e feijão + suco
arroz e feijão + refrigerante
arroz e feijão + água
macarrão + suco
macarrão + refrigerante
macarrão + água
[(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

2 pratos × 3 bebidas = 6 almoços. O repeat=3 combina a mesma sequência com ela mesma três vezes: todas as respostas possíveis de três perguntas de sim (1) ou não (0).

from itertools import product

lanches = ["misto", "pastel", "coxinha"]
sucos = ["laranja", "uva"]

print(lanches)
print(sucos)

from itertools import product

lanches = ["misto", "pastel", "coxinha"]
sucos = ["laranja", "uva"]

combos = list(product(lanches, sucos))
for lanche, suco in combos:
    print(f"{lanche} + suco de {suco}")
print("Total:", len(combos), "combos")

6. Mais três úteis: accumulate, pairwise, batched

from itertools import accumulate, pairwise, batched

movimentos = [100, -30, -20, 50]
print(list(accumulate(movimentos)))

temperaturas = [21, 24, 23, 27]
for antes, depois in pairwise(temperaturas):
    print(f"{antes} -> {depois}: {depois - antes:+}")

convidados = ["Ana", "Bruno", "Carla", "Davi", "Eva"]
for mesa in batched(convidados, 2):
    print(mesa)

$ python3 outros.py
[100, 70, 50, 100]
21 -> 24: +3
24 -> 23: -1
23 -> 27: +4
('Ana', 'Bruno')
('Carla', 'Davi')
('Eva',)

accumulate(movimentos)

"Acumular": a soma até ali, item por item. É o saldo da conta depois de cada movimento: 100, 100 − 30 = 70, 70 − 20 = 50, 50 + 50 = 100.

pairwise(temperaturas)

"Aos pares": cada item com o seguinte, (21, 24), (24, 23), (23, 27). Bom para calcular diferenças. O :+ da f-string mostra o sinal sempre (aula Formatação). Existe desde o Python 3.10.

batched(convidados, 2)

"Em lotes": tuplas de 2 em 2; o último lote fica com o que sobrou. São as mesas de dois lugares da festa. Existe desde o Python 3.12.

Formas diferentes: itertools ou laço na mão?

O laço na mão não está errado, e às vezes é mais fácil de ler para quem está começando. As ferramentas valem a pena quando o laço fica com índices e contas que podem sair errados por um.

Quando dá errado

from itertools import groupby

alunos = ["Ana", "Alice", "Bruno", "Beatriz", "Caio"]
grupos = groupby(alunos, key=lambda nome: nome[0])

guardados = {}
for letra, grupo in grupos:
    guardados[letra] = grupo

for letra, grupo in guardados.items():
    print(letra, list(grupo))

$ python3 grupo_gasto.py
A []
B []
C []

from itertools import groupby

alunos = ["Ana", "Alice", "Bruno", "Beatriz", "Caio"]

guardados = {}
for letra, grupo in groupby(alunos, key=lambda nome: nome[0]):
    guardados[letra] = list(grupo)

print(guardados)

$ python3 grupo_ok.py
{'A': ['Ana', 'Alice'], 'B': ['Bruno', 'Beatriz'], 'C': ['Caio']}

from itertools import count

senhas = count(1)
print(senhas[:5])

$ python3 erro_fatia.py
Traceback (most recent call last):
  File "/home/ana/erro_fatia.py", line 4, in <module>
    print(senhas[:5])
          ~~~~~~^^^^
TypeError: 'itertools.count' object is not subscriptable

from itertools import chain

total = chain(10, 20)
print("Criei a corrente.")
print(list(total))

$ python3 erro_chain.py
Criei a corrente.
Traceback (most recent call last):
  File "/home/ana/erro_chain.py", line 5, in <module>
    print(list(total))
          ~~~~^^^^^^^
TypeError: 'int' object is not iterable

$ python3 print_obj.py
<itertools.combinations object at 0x7f1893929bc0>

Referência rápida

Exercícios

Missão final: as duplas do trabalho

from itertools import combinations

alunos = ["Ana", "Bruno", "Carla", "Davi"]

duplas = []
for a, b in duplas:
    print(a, "e", b)
print("Total de duplas:", len(duplas))

from itertools import combinations

alunos = ["Ana", "Bruno", "Carla", "Davi"]

duplas = list(combinations(alunos, 2))
for a, b in duplas:
    print(a, "e", b)
print("Total de duplas:", len(duplas))

São 6 duplas, de Ana e Bruno até Carla e Davi. Troque combinations por permutations e veja o total dobrar para 12: agora "Ana e Bruno" e "Bruno e Ana" contam separado.

Resumo

• itertools traz iteradores prontos; todos entregam aos poucos e se gastam numa passada.
• count e cycle não têm fim: limite com zip, islice ou break.
• chain emenda sequências (inclusive geradores) sem montar uma lista grande.
• groupby junta só vizinhos: ordene antes pela mesma chave, e transforme cada grupo em lista na mesma volta.
• combinations (sem ordem), permutations (com ordem) e product (um de cada) montam todas as possibilidades; elas crescem depressa.
• accumulate dá somas parciais, pairwise dá vizinhos, batched dá lotes.