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.

Generator Expressions

tl;dr

Generator expressions são mais adequadas do que list comprehensions quando queremos apenas gerar sequências de elementos para iterar sobre, pois geram um elemento de cada vez, ao contrário das list comprehensions que geram uma lista inteira em memória.

Generator Expressions

Frequentemente, você pode economizar memória e processamento em seus programas fazendo:

lista = (x**2 for x in xrange(1, 1000001))

ao invés do costumeiro:

lista = [x**2 for x in xrange(1, 1000001)]

O primeiro trecho de código usa uma generator expression,  que são expressões que retornam um generator (calma aí que já já vou descrever o que é um generator). O segundo trecho de código usa uma list comprehension expression, que retorna uma lista. Para entendermos a diferença entre as duas expressões, primeiro é necessário que saibamos o que é um generator. Se você já sabe, pode pular a próxima seção.

Generators

Não são raras as situações em que precisamos criar uma sequência de elementos, para depois iterar sobre essa sequência realizando alguma operação sobre esses elementos. Por exemplo, queremos gerar uma sequência contendo os n primeiros números primos. Poderíamos criar uma funçãozinha que gera uma lista com esses valores:

def primos(n):
    lista = []
    i = 0
    while len(lista) < n:
        if eh_primo(i):
            lista.append(i)
        i = i + 1
    return lista

E quando precisarmos usar tal lista, simplesmente chamamos a função primos(). Assim, podemos iterar sobre o resultado da chamada a essa função:

for i in primos(1000):
    # faça algo com o i
    print i,

O problema é que a função primos() gera uma lista de n elementos em memória — e isso pode ser um tanto quanto caro, dependendo do valor de n. Como estamos apenas querendo iterar sobre um conjunto de elementos (ou seja, usar cada um de uma vez), poderíamos usar um generator. Para entender o que é um generator, vamos primeiro reescrever a função anterior:

def primos_gen(n):
    i = 1
    count = 0
    while count < n:
        if eh_primo(i):
            count = count + 1
            yield i
        i = i + 1

Repare na palavra-chave yield. Em Python, toda função que possui em seu corpo a instrução yield, retorna um generator quando for chamada. O yield pode ser lido como um pause, que retorna um objeto do tipo generator. Veja:

>>> x = primos_gen(100)
>>> print type(x)
<type 'generator'>

Sendo um objeto do tipo generator, x possui um método next(), que irá retornar um elemento apenas. Veja que a cada chamada ao next(), o generator x retorna o valor seguinte da sequência que está sendo gerada:

>>> x.next()
1
>>> x.next()
2
>>> x.next()
3
>>> x.next()
5
>>> x.next()
7
>>> x.next()
11

E assim sucessivamente… O que deve ser entendido sobre o comando yield é que ele é parecido com o return, pois acaba retornando um elemento para quem chamar a função next(). Porém, a execução da função generator fica “pausada” após o yield, sem perder o contexto atual, como ocorreria no caso de um return. Assim, quando chamarmos novamente o método next() do objeto, a execução irá continuar na linha seguinte ao yield, não reiniciando a execução da função desde o seu início. Ou seja, a função volta a ser executada do ponto onde estava parada. Legal, né?

Isso permite que façamos uma iteração sobre uma sequência de elementos sem ter que gerar toda essa sequência em memória antecipadamente. Assim, geramos elemento por elemento e o retornamos, e, quando quisermos outro elemento, a execução da função geradora será retomada de onde parou na execução anterior. Caso ainda não tenha entendido o yield, observe e teste o seguinte código (emprestado daqui):

def func():
    yield 1
    yield 2
    yield 3
x = func()
print x.next()
print x.next()
print x.next()

Se você testar o código, verá que quando executamos a primeira chamada à função x.next(), teremos como valor de retorno o número 1. Quando executamos a segunda chamada à x.next(), teremos como retorno o valor 2. Por fim, quando executamos a terceira chamada à x.next(), teremos como resposta o inteiro 3.

Mas, ao invés de invocar explicitamente o método next() de x, poderíamos usar um laço for (repare que não estamos chamando o next()). Veja:

def func():
    yield 1
    yield 2
    yield 3

x = func()
for i in x:
    print i,

Execute o código acima e verá que teremos a seguinte saída:

1 2 3

Parabéns, você acaba de aprender a “mágica” que o laço for usa para percorrer sequências de elementos! Na verdade, o for espera que o objeto a ser percorrido retorne para ele um iterator, para que ele possa “conversar” com tal objeto, chamando o método next() a cada iteração para obter o próximo valor.

Enfim, generators são objetos que retornam objetos e que mantém o estado da função geradora entre a geração de um objeto e outro.

Generator Expressions

Como comentei anteriormente, muitas vezes usamos list comprehensions quando na verdade elas não são a melhor alternativa. Quando usamos list comprehensions, a lista inteira é gerada de uma vez e armazenada na memória. Por exemplo, se quisermos gerar uma lista contendo os quadrados de todos os números inteiros de 1 a 1.000.000, podemos fazer o seguinte:

>>> lista = [x**2 for x in xrange(1, 1000001)]

Ao usarmos uma list comprehension, estamos gerando uma lista inteira na memória e armazenando-a em uma variável chamada lista. Tudo bem, às vezes é mesmo necessário manter essa lista na memória para usar seu conteúdo posteriormente, mas muitas vezes fazendo mau uso da ferramenta e usamos esse tipo de expressão quando não precisamos da lista inteira de uma vez só.

Vamos seguir o raciocínio utilizando um exemplo. Estudando o módulo random, desejamos verificar, em um conjunto de 1000000 (um milhão) de números aleatórios gerados através da função random.randint(), onde limitamos os números gerados à faixa de 0 a 1000000, se alguma vez o valor 0 é gerado pela função. Usando uma list comprehension e a função any(), poderíamos verificar se algum dentre os 1000000 números gerados é igual a 0. (Perceba que o resultado vai variar de execução para execução, dependendo dos números gerados.)

>>> any([random.randint(0, 1000000) == 0 for x in xrange(0, 1000000)])
True
>>> any([random.randint(0, 1000000) == 0 for x in xrange(0, 1000000)])
False

O código acima basicamente executa os seguintes passos:

  1. A cada volta do loop é gerado um número aleatório entre 0 e 1000000.
  2. Para cada número gerado, verificamos se ele é igual a 0 e armazenamos o booleano correspondente a tal verificação em uma lista (assim, temos uma lista de 1000000 de booleanos, por exemplo: [False, True, False, False, …, False]).
  3. Após gerada a lista, usamos a função any(), que verifica se algum dos valores armazenados na lista é True.

Nosso código tem um sério problema. Antes de dizer qual é, pergunto a vocês: “Precisamos realmente armazenar na memória uma lista de 1000000 de elementos para fazer a operação acima?”. A resposta é: “Nesse caso, não!”. A função any() recebe um objeto iterável como parâmetro e irá chamar o método next() do iterator retornado por esse objeto a cada iteração necessária. Poderíamos usar uma generator expression nesse caso. Esse tipo de expressão se comporta de forma semelhante a uma função do tipo generator. Ao contrário de quando iteramos sobre uma list comprehension, onde geramos uma lista inteira antes de percorrê-la, com uma generator expression, os elementos são gerados “sob demanda”. Podemos escrever o mesmo exemplo acima utilizando a sintaxe das generator expressions (muito semelhante à list comprehensions, trocando o [] por ()):

>>> any((random.randint(0, 1000000) == 0 for x in xrange(0, 1000000)))
True
>>> any((random.randint(0, 1000000) == 0 for x in xrange(0, 1000000)))
False

A função any() itera sobre o objeto gerado pela generator expression. Temos então a geração de um elemento a cada iteração. Dessa forma, estamos economizando memória, pois os elementos são gerados e, assim que não são mais necessários, podem ser descartados.

As generator expressions estão disponíveis a partir da versão 2.4 do Python.

List comprehensions são do mal?

Claro que não. Em muitos casos, list comprehensions são exatamente o que precisamos. Elas geram listas, que permitem acesso aleatório a elementos dela (ou seja, podemos acessar o n-ésimo elemento dela sem antes passar pelos n-1 elementos anteriores). Também permite que façamos uso do operador de slicing, dentre outras vantagens das listas.

Mais informações

Veja mais na PEP que propôs a criação de generator expressions. Outro material muito bom são os slides da palestra que o Luciano Ramalho fez no FISL/2012 sobre esse assunto.

map(), reduce(), filter() e lambda

map()

map() é uma função builtin de Python, isto é, uma função que é implementada diretamente no interpretador Python, podendo ser utilizada sem a importação de um módulo específico. Essa função, em Python, serve para aplicarmos uma função a cada elemento de uma lista, retornando uma nova lista contendo os elementos resultantes da aplicação da função. Considere o exemplo abaixo:

>>> import math
>>> lista1 = [1, 4, 9, 16, 25]
>>> lista2 = map(math.sqrt, lista1)
>>> print lista2
[1.0, 2.0, 3.0, 4.0, 5.0]

Ao chamar a função map(math.sqrt, lista1), estamos solicitando ao interpretador para que execute a função math.sqrt (raiz quadrada – square root) usando como entrada cada um dos elementos de lista1, e inserindo o resultado na lista retornada como resultado da função map().

É uma forma bem interessante e expressiva de denotar a aplicação de uma função a cada elemento de uma lista (ou sequência). Mas, podemos facilmente substituir uma chamada a map() por list comprehensions. O código recém listado poderia ser substituído por:

>>> lista2 = [math.sqrt(x) for x in lista1]
>>> print lista2
[1.0, 2.0, 3.0, 4.0, 5.0]

O código acima produz o mesmo resultado que map(), pois, para cada elemento de lista1, executa a função math.sqrt e inclui o resultado dessa execução na lista de retorno.

O fato de a função map() ser tão facilmente substituída pelo uso de comprehensions, já criou até mesmo algumas discussões sobre manter ou não map() entre as funções builtin do Python 3000 [1].

reduce()

reduce() é outra função builtin do Python (deixou de ser builtin e passou a estar disponível no módulo functools a partir da versão 3000). Sua utilidade está na aplicação de uma função a todos os valores do conjunto, de forma a agregá-los todos em um único valor. Por exemplo, para aplicar a operação de soma a todos os elementos de uma sequência, de forma que o resultado final seja a soma de todos esses elementos, poderíamos fazer o seguinte:

>>> import operator #necessário para obter a função de soma
>>> valores = [1, 2, 3, 4, 5]
>>> soma = reduce(operator.add, valores)
>>> print soma
15

É claro que, para realizar a soma de todos os elementos de uma sequência, é muito mais claro utilizarmos a função sum():

>>> print sum(valores)
15

Como falei anteriormente, reduce() foi retirada do conjunto de builtins de Python, em parte devido à obscuridade que pode acabar gerando [1].

filter()

Como o próprio nome já deixa claro, filter() filtra os elementos de uma sequência. O processo de filtragem é definido a partir de uma função que o programador passa como primeiro argumento da função. Assim, filter() só “deixa passar” para a sequência resultante aqueles elementos para os quais a chamada da função que o usuário passou retornar True. Vejamos um exemplo:

>>> def maior_que_zero(x):
...     return x > 0
...
>>> valores = [10, 4, -1, 3, 5, -9, -11]
>>> print filter(maior_que_zero, valores)
[10, 4, 3, 5]

No exemplo acima, filter() chamou a função maior_que_zero para cada um dos elementos contidos em valores. Se a função retornar True, o elemento é inserido na lista de resultado. Caso contrário, não o é. Ou seja, é feita a filtragem e só passam aqueles elementos que são maiores que zero.

Assim, como no exemplo da builtin map(), aqui também podemos escrever com facilidade uma comprehension com a mesma funcionalidade:

>>> print [x for x in valores if x > 0]
[10, 4, 3, 5]

Devido a essa fácil substituição, filter() também já esteve na mira para ser retirada do conjunto de builtins, embora tenha acabado permanecendo.

lambda

No exemplo da função filter(), tivemos que definir uma nova função (chamada maior_que_zero) para usar somente dentro da função filter(), sendo chamada uma vez para cada elemento. Ao invés de definir uma nova função dessa forma, poderíamos definir uma função válida somente enquanto durar a execução do filter. Não é necessário nem dar um nome a tal função, sendo portanto chamada de função anônima ou função lambda. Considere o exemplo abaixo:

>>> valores = [10, 4, -1, 3, 5, -9, -11]
>>> print filter(lambda x: x > 0, valores)
[10, 4, 3, 5]

Definimos uma função anônima (portanto, não tem nome), que recebe uma entrada (a variável x) e retorna o resultado da operação x > 0, True ou False.

Poderíamos também usar uma função lambda no exemplo da função reduce():

>>> valores = [1, 2, 3, 4, 5]
>>> soma = reduce(lambda x, y: x + y, valores)
>>> print soma
15

No código acima, definimos uma função anônima que recebe duas entradas e retorna a soma dessas entradas.

[1] Guido Van Rossum. The fate of reduce() in Python 3000