each_cons — percorrendo sequências em N elementos por vez

Recentemente descobri o each_cons, um método interessante da API de Ruby no mixin Enumerable. Esse método permite você percorrer uma sequência de tantos em tantos elementos por vez.

Por exemplo, se você tem uma lista [1, 3, 5, 7, 9] você pode usar o each_cons para percorrer de 2 em 2 elementos:

(1, 3), (3, 5), (5, 7), (7, 9)

ou de 3 em 3 elementos de cada vez:

(1, 3, 5), (3, 5, 7), (5, 7, 9)

Ele ajuda a criar operações que envolvem uma espécie de “janela deslizante” (sliding window), um conjunto de elementos da sequência sobre o qual podemos inferir alguma informação.

Essa abstração se revela bem útil na hora de escrever certos algoritmos e por isso é interessante tê-la como carta na manga.

Implementando em Python

Python não tem um método ou função equivalente ao each_cons na API padrão, mas encontrei algumas sugestões de implementação nas respostas a uma pergunta no StackOverflow. Eis uma delas, que funciona para listas e tuplas:

import itertools
def each_cons(xs, n):
    return itertools.izip(*(xs[i:] for i in xrange(n)))

Para entender como esse código funciona, vamos por partes. Acompanhe:

>>> xs = [1, 2, 3, 4, 5]
>>> xs[0:] # primeiro slice
[1, 2, 3, 4, 5]
>>> xs[1:] # segundo slice
[2, 3, 4, 5]
>>> xs[2:] # terceiro slice
[3, 4, 5]
>>> zip(xs[0:], xs[1:], xs[2:]) # zip dos slices
[(1, 2, 3), (2, 3, 4), (3, 4, 5)]

Os passos acima delineiam o algoritmo dessa implementação do each_cons. Basicamente, obtemos slices da lista na quantidade de elementos que queremos na nossa janela deslizante, e aplicamos zip() neles.

O código mostrado usa uma generator expression para gerar os slices, o truque da estrela nos argumentos para passar os argumentos pro zip e a função izip do módulo itertools que é o equivalente da builtin zip() para generators.

Exemplos de utilidade

1) Encontrar os pares de números com distância entre si maior que 1:

>>> numeros = [1, 2, 3, 5, 9, 10]
>>> [(a, b) for a, b in each_cons(numeros, 2) if b - a > 1]
[(3, 5), (5, 9)]

2) Descobrir se os números de uma lista formam uma sequência:

>>> all([a + 1 == b for a, b in
                        each_cons([1, 2, 3, 4, 5], 2)])
True
>>> all([a + 1 == b for a, b in
                        each_cons([1, 3, 5, 7], 2)])
False

3) Calcular as médias das vendas de cada trimestre:

>>> totais_mensais = [123.45, 54.3, 428, 144.2, 245.45]
>>> [float(a+b+c)/3 for a, b, c in
                        each_cons(totais_mensais, 3)]
[201.91666666666666, 208.83333333333334, 272.55]

Esse tipo de média é conhecido por média móvel, é comum em aplicações financeiras e também é útil para amaciar as linhas de um gráfico.

4) Percorrer as linhas de um arquivo procurando por duplicatas:

>>> linhas = open('arquivo.txt').readlines()
>>> [(num, a) for num,(a,b) in
                  enumerate(each_cons(linhas, 2), 1)
                  if a == b]
[(3, 'tres\n'), (5, 'quatro\n')]

Para um arquivo.txt contendo o texto:

um
dois
tres
tres
quatro
quatro
cinco

Generalizando para generators

A função each_cons que apresentamos acima, apesar de funcionar bem para listas e tuplas, não funciona para generators nem para objetos xrange(). Veja o erro que acontece quando tentamos rodar com um generator:

>>> each_cons((a for a in xrange(10)), 2) # testando com um generator
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in each_cons
TypeError: type object argument after * must be a sequence, not generator

É interessante termos uma solução que funcione também para generators, para o caso de estarmos lidando com um grande volume de dados vindo de um banco de dados, por exemplo.

Para obter isso, basta usar mais algumas funções do módulo itertools na brincadeira:

import itertools
def each_cons(xs, n):
    return itertools.izip(*(itertools.islice(g, i, None)
                          for i,g in
                          enumerate(itertools.tee(xs, n))))

A implementação mostrada acima funciona com sequências (listas, tuplas, strings), generators e objetos xrange() também.

Um comentário sobre “each_cons — percorrendo sequências em N elementos por vez

  1. Pingback: Como funcionam as listas de Python? | Python Help

Deixe uma resposta

Preencha os seus dados abaixo ou clique em um ícone para log in:

Logotipo do WordPress.com

Você está comentando utilizando sua conta WordPress.com. Sair / Alterar )

Imagem do Twitter

Você está comentando utilizando sua conta Twitter. Sair / Alterar )

Foto do Facebook

Você está comentando utilizando sua conta Facebook. Sair / Alterar )

Foto do Google+

Você está comentando utilizando sua conta Google+. Sair / Alterar )

Conectando a %s