Curso de Python / Contar, somar e agrupar com collections e itertools

Use collections para contar ocorrências ou acumular valores e itertools para combinar fluxos de iteração. Escolha a ferramenta pelo significado do resultado, pela ordem dos dados e pelo consumo dos iteradores, não apenas pela quantidade de linhas de código.

Separar contagem e soma por categoria

Counter conta quantas vezes cada item aparece. No exemplo desta aula, há duas despesas na categoria food, mas seu total é 11.500 unidades. Contar duas ocorrências e somar seus valores são operações diferentes; trocar uma pela outra produziria um relatório com outro significado.

defaultdict(int) permite acumular valores sem inicializar cada categoria manualmente. Ao consultar uma chave inexistente com colchetes, a fábrica int cria o zero e o dicionário passa a conter aquela chave. Por isso, uma simples consulta pode modificar o contêiner. Já totals.get("unknown") devolve None sem chamar a fábrica nem inserir a chave.

Se a leitura não deve criar categorias, considere um dicionário comum com get e atribuição explícita na etapa de acumulação. A documentação de collections.defaultdict descreve o comportamento da fábrica.

Executar contagens, totais e encadeamento

Salve este programa completo como collections_demo.py, em UTF-8, em uma pasta de prática. Não é necessário instalar pacotes: os dois módulos fazem parte da biblioteca padrão. Execute python collections_demo.py no ambiente virtual do projeto. Fora dele no macOS ou Ubuntu, use python3 se esse for o comando do interpretador.

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

Resultado esperado:

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

A expressão geradora passada a Counter entrega somente as categorias; o sublinhado indica que o valor não é usado nessa contagem. O laço seguinte usa os valores para preencher totals. As chamadas a sorted(...items()) deixam a apresentação por chave previsível. Elas não fazem o Counter somar quantias.

chain([1, 2], [3]) percorre a primeira coleção e depois a segunda. A chamada a list materializa os itens para imprimi-los. Na última parte, groupby recebe as linhas ordenadas: com estas tuplas, a ordenação padrão começa pela categoria e depois compara o valor. Todos os registros da mesma categoria ficam próximos.

Agrupar chaves consecutivas com groupby

itertools.groupby agrupa trechos consecutivos com a mesma chave. Ele não procura automaticamente todos os registros iguais espalhados pelo conjunto, como um agrupamento SQL. Se a sequência for food, bus, food, haverá três grupos. Ordenar pela categoria antes de agrupar reúne os dois registros food.

A função key extrai o valor usado na comparação. lambda row: row[0] é uma função anônima curta que faz o mesmo que a função category abaixo. Para tornar a intenção clara, este segundo programa usa a mesma função tanto em sorted quanto em groupby.

Salve o código completo como groups.py na mesma pasta e execute 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]))

Resultado esperado:

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

Os rótulos original e sorted significam, respectivamente, ordem original e ordem classificada. A primeira execução imprime dois grupos food; a segunda imprime um único total food de 11.500. A última linha confirma que a função nomeada e a expressão lambda extraem o mesmo valor deste registro.

O iterador de cada grupo compartilha a origem com o agrupador. Consuma o grupo antes de avançar ao próximo: guardar apenas o objeto iterador e tentar lê-lo depois pode resultar em um grupo vazio. Se precisar reutilizar os detalhes, materialize list(group) enquanto o grupo estiver ativo; se só precisa do total, armazene a soma calculada. Veja as condições de itertools.groupby e a explicação sobre consumo de iteradores e geradores.

Avaliar memória e conferir a escolha

chain evita criar previamente uma lista concatenada, mas envolver o resultado inteiro em list volta a armazenar todos os itens. islice permite consumir uma parte delimitada de um iterador; use um limite explícito ao trabalhar com uma fonte infinita. Não tente converter uma fonte infinita inteira em lista.

A etapa sorted(rows, key=category) cria uma lista e precisa consumir a entrada para ordená-la. Portanto, o exemplo não é um processamento de memória constante só porque usa groupby. Se os dados já chegarem ordenados pela categoria, você pode dispensar essa etapa; caso contrário, uma acumulação por chave pode atender melhor à necessidade de totais.

1. No primeiro arquivo, remova somente sorted da entrada de groupby, passando rows diretamente. As contagens e os totais iniciais continuam iguais, mas a parte final mostra food 3500, bus 2500 e food 8000.
2. Restaure a ordenação e confirme os totais originais.
3. Altere a despesa food de 8000 para 9000. Sua contagem permanece 2, enquanto a soma sobe para 12.500. Restaure o valor.
4. Antes de imprimir os totais, consulte uma chave ausente com colchetes e depois com get, usando nomes diferentes. Confira quais chaves passaram a existir e restaure o arquivo. Essa diferença ajuda a decidir se a criação automática é desejável.