Curso de Python / Agrupar y combinar datos con collections e itertools

collections ofrece contenedores especializados para contar y acumular datos, mientras que itertools combina recorridos de elementos. La elección depende de qué significa el resultado y de cómo se consume la entrada, no solo de escribir menos líneas.

Distinguir frecuencias y totales

Counter cuenta apariciones: dos registros de comida representan dos apariciones, no el coste total de ambos. defaultdict(int) facilita acumular importes porque crea un cero al acceder mediante corchetes a una clave que todavía no existe.

Ese acceso también inserta la clave. Si solo quieres consultar sin modificar el diccionario, un diccionario ordinario con get puede expresar mejor tu intención. Conviene repasar las reglas de claves y valores antes de atribuir a estas herramientas comportamientos que no tienen.

Ejecutar recuentos y agrupaciones

Guarda el programa como collections_demo.py y ejecútalo con python collections_demo.py. Ambos módulos vienen con Python. Las categorías del ejemplo permanecen en inglés como etiquetas de datos para poder comparar los resultados entre lecciones.

from collections import Counter, defaultdict
from itertools import chain, groupby

rows = [("food", 3500), ("bus", 2500), ("food", 8000)]
counts = Counter(category for category, _ in rows)
totals = defaultdict(int)
for category, amount in rows:
    totals[category] += amount
print(sorted(counts.items()))
print(sorted(totals.items()))
print(list(chain([1, 2], [3])))
for category, group in groupby(sorted(rows), key=lambda row: row[0]):
    print(category, sum(amount for _, amount in group))

[('bus', 1), ('food', 2)]
[('bus', 2500), ('food', 11500)]
[1, 2, 3]
bus 2500
food 11500

La salida distingue dos registros de comida y su suma de 11500. chain recorre primero una lista y después la otra. En la última parte, los registros se ordenan antes de agruparlos para que las categorías iguales queden juntas.

Comprender qué agrupa realmente groupby

itertools.groupby reúne claves iguales consecutivas. No busca automáticamente todos los registros de una categoría dispersos en la entrada como una agrupación SQL. key recibe una función que extrae la clave; lambda row: row[0] es una función breve que devuelve la categoría.

El siguiente archivo independiente, groups.py, compara los mismos registros en su orden original y ordenados por esa clave. Ejecútalo con python groups.py.

from itertools import groupby
rows = [("food", 3500), ("bus", 2500), ("food", 8000)]
def category(row):
    return row[0]
for label, data in [("original", rows), ("sorted", sorted(rows, key=category))]:
    print(label)
    for name, group in groupby(data, key=category):
        print(name, sum(amount for _, amount in group))
print(category(rows[0]) == (lambda row: row[0])(rows[0]))

original
food 3500
bus 2500
food 8000
sorted
bus 2500
food 11500
True

La entrada original produce dos grupos de comida separados por el autobús. La entrada ordenada produce uno solo. La función category y la expresión lambda obtienen la misma clave, como confirma la última línea.

Los iteradores de cada grupo comparten el flujo subyacente. Consúmelos antes de avanzar al siguiente grupo o guarda los datos necesarios en una lista; no supongas que podrás volver a leer un grupo antiguo después.

Conservar solo los resultados necesarios

chain no necesita construir previamente la concatenación completa. Sin embargo, envolver el resultado en list vuelve a almacenarlo entero. Para tomar una parte de un iterador puedes utilizar islice; un iterador infinito necesita un límite o una condición de parada. El modelo de generadores e iteradores explica estas consecuencias.

Elimina temporalmente la ordenación del primer groupby, observa los dos grupos de comida y restáurala. Decide primero si necesitas frecuencias, importes acumulados o grupos consecutivos. Las referencias de defaultdict y groupby detallan sus reglas de creación y consumo.