Brincando de arvore binária com Python

De repente me deu vontade de fazer uma árvore binária em python. Não sei por que, nem para quê, mas deu vontade.

Usei um exemplo que creio ter visto no livro “C A linguagem de Programação” de Brian Kernighan e Dennis Ritchie : contagem de palavras em um texto.

O livro fala de árvores binárias como alternativa, neste problema de contagens de palavras, à pesquisa binária. Numa pesquisa binária, pressupomos que a lista de valores está previamente ordenada. Isso não deve ocorrer com as palavras de um texto qualquer. Não sabemos nem mesmo a quantidade de palavras que podem estar contidas nesse texto, de modo que pode não ser muito eficiente rodar previamente um algoritmo de ordenação. A salvação então é a tal árvore binária.

No meu código abaixo, criei uma classe (arvbin) para representar os nós da árvore. Os atributos da classe são: a palavra, o contador de ocorrências da palavra, um apontador para a palavra menor e outro para a maior:

class arvbin:
    def __init__(self, Valor):
        self.Menor=None  #apontador para a palavra menor
        self.Maior=None  #apontador para a palavra maior
        self.conta=1     #conttador de ocorrências
        self.valor=Valor #a palavra

O método add percorre a árvore até achar um lugar apropriado para a palavra a ser adicionada. Em cada nó, ele compara a nova palavra à palavra do nó. Se forem iguais, simplesmente incrementa o contador e retorna; se for maior, verifica se ja existe um nó maior no apontador sef.Maior; não tendo, adiciona um novo nó ao apontador “self.Maior”, com a palavra nova. Se já existe o nó maior, repete a operação, agora já no nó maior. Se a nova palavra for menor, dá-se a mesma operação, só que a partir do apontador self.Menor.

O método busca percorre a árvore buscando um nó que contenha determinada palavra.

O método ordena percorre recursivamente a árvore e retorna uma lista ordenada com as palavras.

#! /usr/bin/env python
# -.- encoding: utf-8 -.-

class arvbin:
    def __init__(self, Valor):
        self.Menor=None
        self.Maior=None
        self.conta=1
        self.valor=Valor

    def add(self,valor):
        if valor>self.valor:
            self.addmaior(valor)
        elif valor<self.valor:
            self.addmenor(valor)
        else:
            self.conta=self.conta+1

    def addmenor(self,valor):
        if self.Menor:
            self.Menor.add(valor)
        else:
            self.Menor=arvbin(valor)

    def addmaior(self,valor):
        if self.Maior:
            self.Maior.add(valor)
        else:
            self.Maior=arvbin(valor)

    def busca(self,valor):
        if valor>self.valor:
            if self.Maior:
                return self.Maior.busca(valor)
            else :
                return None
        elif valor<self.valor:
            if self.Menor:
                return self.Menor.busca(valor)
            else :
                return None
        else:
            return self

    def ordena(self):
        listamenor=[]
        listamaior=[]
        if self.Menor:
            listamenor=self.Menor.ordena()
        if self.Maior:
            listamaior=self.Maior.ordena()
        return listamenor + [ (self.valor, self.conta) ] + listamaior 

Para testar meu código, adicionei o trecho abaixo. Este código supõe que já existe um arquivo “/tmp/texto” contendo as palavras da nossa brincadeira. Dá para criar um arquivo assim simplesmente copiando um grande texto na web.

if __name__ == '__main__':
    texto=open('/tmp/texto','r' )
    a=texto.read()

    remover=[ '.',',',';',':','?','!','(',')','\'','"' ]
    for remova in remover:
        a=a.replace(remova,'' )

    a=a.replace('\n',' ' )

    a=a.split()
    raiz=arvbin(a[-1] )
    a.pop()
    for palavra in a:
        raiz.add(palavra.upper())
    for palavra in raiz.ordena():
        print palavra[1], palavra[0]

    achei = raiz.busca('não' )
    if achei:
        print achei.valor, achei.conta