Módulo itertools de Python

El módulo itertools de Python es un conjunto de herramientas de la biblioteca estándar con bloques de construcción rápidos y eficientes en memoria para trabajar con iteradores. Cada función de itertools devuelve un iterador — produce valores bajo demanda en lugar de construir una lista en memoria — lo que hace que el módulo sea ideal para grandes conjuntos de datos, secuencias infinitas y canalizaciones de datos componibles.

Este capítulo cubre las tres categorías de funciones de itertools: iteradores infinitos (count, cycle, repeat), iteradores combinatorios (product, permutations, combinations, combinations_with_replacement), e iteradores terminantes (chain, islice, groupby, compress, filterfalse, takewhile, dropwhile, starmap, zip_longest, accumulate, pairwise).

No se requiere instalación — itertools viene incluido en cada instalación de Python 3:

import itertools

¿Por qué usar itertools?

Considera leer los primeros 10 múltiplos de un número. Sin itertools necesitas una lista o un contador manual. Con itertools.count e itertools.islice la intención queda inmediatamente clara y el uso de memoria 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]

La filosofía de itertools: construir una pieza pequeña y correcta, y luego componerla con otras. Encadenar dos funciones de itertools es más rápido y menos propenso a errores que escribir el bucle equivalente a mano.

Iteradores Infinitos

Estos iteradores producen valores indefinidamente. Combínalos siempre con islice, un for … break, u otro mecanismo terminante para evitar un bucle infinito.

count(start=0, step=1)

count produce una secuencia de números con espaciado uniforme. Es esencialmente range sin límite superior y con soporte para flotantes y pasos 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 es útil cuando necesitas numerar elementos de un iterable sin saber cuántos hay de antemano — el idioma de enumerate pero con inicio y paso personalizados.

cycle(iterable)

cycle repite los elementos de cualquier iterable 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áctico — asignar elementos a equipos en 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 produce el mismo objeto times veces (o indefinidamente si se omite times).

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]

El patrón map(pow, range(1, 6), repeat(2)) es un idioma común para suministrar un segundo argumento constante a una función de dos argumentos.

Iteradores Combinatorios

Estos iteradores producen todas las combinaciones, permutaciones o productos cruzados de un iterable de entrada. Son esenciales para búsquedas por fuerza bruta, generación de casos de prueba y problemas de combinatoria.

product(*iterables, repeat=1)

product calcula el producto cartesiano — todas las combinaciones ordenadas donde se toma un elemento de cada iterable. Es equivalente a bucles for anidados.

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

Usa repeat para calcular el producto de un iterable consigo mismo varias veces:

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

Advertencia: product materializa los iterables de entrada en memoria (para permitir múltiples pasadas), así que no pases iteradores enormes como entrada.

permutations(iterable, r=None)

permutations produce todas las ordenaciones de r elementos tomados del input. Cuando se omite r, se usan todos los elementos.

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)

El orden importa en las permutaciones — ('A', 'B') y ('B', 'A') son resultados distintos.

combinations(iterable, r)

combinations produce todas las selecciones no ordenadas de r elementos. A diferencia de permutations, el orden no importa — cada subconjunto aparece solo una 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)

Un caso de uso común — verificar todos los pares de elementos por una propiedad:

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, pero permite que cada elemento aparezca más de una vez en una selección.

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

Esto es útil para generar todas las tiradas posibles de dados, secuencias de lanzamiento de monedas, o elecciones de caracteres para contraseñas.

Funciones combinatorias de un vistazo

Iteradores Terminantes

Los iteradores terminantes procesan una entrada finita y se detienen cuando esa entrada se agota.

chain(*iterables)

chain trata varios iterables como una única secuencia continua sin construir una nueva 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 acepta un único iterable de iterables — útil cuando no se conoce el número de secuencias de antemano:

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 es una alternativa rápida y eficiente en memoria a [item for sublist in nested for item in sublist].

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

islice corta cualquier iterador — incluidos los infinitos — sin materializarlo. Los argumentos imitan la notación slice de Python pero solo aceptan enteros no 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 no admite índices negativos ni pasos negativos (a diferencia del corte de listas regular).

groupby(iterable, key=None)

groupby agrupa elementos consecutivos que comparten el mismo valor de clave. Devuelve 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']

Advertencia importante: groupby solo agrupa elementos iguales consecutivos. Si tus datos no están pre-ordenados por la clave, elementos similares en posiciones distintas forman grupos separados (como se muestra arriba — 'cherry' inicia un nuevo grupo 'fruit' en lugar de unirse al primero). Ordena siempre por la clave antes de llamar a 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']

También ten en cuenta que el iterador de grupo se invalida una vez que avanzas a la siguiente clave — consume cada grupo antes de llamar a next() en el iterador externo.

compress(data, selectors)

compress filtra data conservando solo los elementos cuyo valor de selector correspondiente es verdadero.

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 es equivalente a [d for d, s in zip(data, selectors) if s] pero es más rápido y evita la lista intermedia.

filterfalse(predicate, iterable)

filterfalse es el complemento del filter incorporado — produce elementos para los cuales el predicado devuelve 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 produce elementos mientras el predicado sea True, luego se detiene inmediatamente — incluso si los elementos posteriores satisfacerían el 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 es el espejo de takewhile: omite elementos mientras el predicado sea True, luego produce todos los elementos restantes (incluidos aquellos donde el predicado volvería a ser True).

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 y dropwhile son útiles para procesar archivos de log o flujos donde deseas omitir una sección de encabezado o detenerte en una línea centinela.

starmap(function, iterable)

starmap aplica una función a cada elemento de un iterable, desempaquetando el elemento como argumentos posicionales. Es el equivalente de map pero para iterables de tuplas.

import itertools

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

Compara con map(pow, [2, 4, 10], [3, 2, 3]) — starmap funciona cuando tus argumentos ya están agrupados como tuplas.

zip_longest(*iterables, fillvalue=None)

El zip incorporado se detiene en el iterable más corto. zip_longest rellena los iterables más cortos con fillvalue para que todos los iterables se consuman completamente.

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 totales acumulados (o cualquier otra agregación acumulativa). Por defecto suma, pero puedes pasar cualquier función de dos 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]

El parámetro initial (Python 3.8+) antepone un valor semilla antes del primer elemento:

import itertools

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

pairwise(iterable)

pairwise (Python 3.10+) produce pares consecutivos superpuestos del iterable.

import itertools

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

Esto es útil para calcular diferencias entre valores consecutivos, o para lógica de ventana deslizante donde el tamaño de la ventana es exactamente 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 de Python 3.10, el equivalente era zip(data, data[1:]) (funciona para secuencias) o un enfoque manual basado en tee (funciona para iteradores arbitrarios).

Componiendo itertools en canalizaciones

El verdadero poder de itertools surge cuando combinas funciones. Dado que cada función devuelve un iterador, puedes encadenarlas sin listas intermedias.

Ejemplo: las 3 palabras más frecuentes en un 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)]

Ejemplo: dividir un iterable en fragmentos de tamaño fijo

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,)

Python 3.12 incluye itertools.batched de forma nativa, por lo que puedes reemplazar el helper anterior con itertools.batched(data, 3) en Python moderno.

Referencia rápida

Para los conceptos de evaluación perezosa detrás de itertools, consulta Generadores de Python e Iteradores de Python. Para helpers de estilo funcional que complementan itertools, consulta Funciones Lambda de Python y el Módulo collections de Python.

Práctica