2ª Contribuição de Pesquisa

Árvore AVL

from tkinter import *
import math

class NoArv:  
    def __init__(self, dado):
        self.setFilhos(None, None)
        self.dado = dado

    def setFilhos(self, esquerdo, direito):
        self.esquerdo = esquerdo
        self.direito = direito

    def profArvore(self):
        profEsquerda = 0
        if self.esquerdo:
            profEsquerda = self.esquerdo.profArvore()
        profDireita = 0
        if self.direito:
            profDireita = self.direito.profArvore()
        return 1 + max(profEsquerda, profDireita)
    
    def balanco(self):
        profEsquerda = 0
        if self.esquerdo:
            profEsquerda = self.esquerdo.profArvore()
        profDireita = 0
        if self.direito:
            profDireita = self.direito.profArvore()
        return profEsquerda - profDireita
    
class Balanceamento:
    def __init__(self):
        self.raiz = None

    def criarNo(self, dado): 
        return NoArv(dado)

    def profundidadeMax(self, raiz):
        if raiz == None:
            return 0
        return raiz.profArvore()

    def rotDir(self, raiz):
        if raiz :
            raiz.dado, raiz.esquerdo.dado = raiz.esquerdo.dado, raiz.dado
            old_direito = raiz.direito
            raiz.setFilhos(raiz.esquerdo.esquerdo, raiz.esquerdo)
            raiz.direito.setFilhos(raiz.direito.direito, old_direito)
            
    def rotEsq(self, raiz):
        if raiz :
            raiz.dado, raiz.direito.dado = raiz.direito.dado, raiz.dado
            old_esquerdo = raiz.esquerdo
            raiz.setFilhos(raiz.direito, raiz.direito.direito)
            raiz.esquerdo.setFilhos(old_esquerdo, raiz.esquerdo.esquerdo)

    def rotDirEsq(self, raiz):
        if raiz :
            self.rotDir(raiz.direito)
            self.rotEsq(raiz)

    def rotEsqDir(self, raiz):
        if raiz :
            self.rotEsq(raiz.esquerdo)
            self.rotDir(raiz)

    def balancear(self, raiz):
        if raiz :
            balanco = raiz.balanco()

            if balanco > 1:
                if raiz.esquerdo.balanco() > 0:
                    self.rotDir(raiz)
                else:
                    self.rotEsqDir(raiz)
            elif balanco < -1:
                if raiz.direito.balanco() < 0:
                    self.rotEsq(raiz)
                else:
                    self.rotDirEsq(raiz)

    def insValor(self, raiz, dado):
        if raiz == None:
            return self.criarNo(dado)
        else:
            if dado <= raiz.dado:
                raiz.esquerdo = self.insValor(raiz.esquerdo, dado)
            else:
                raiz.direito = self.insValor(raiz.direito, dado)
                
            self.balancear(raiz)

        return raiz

class Grafico:
    def __init__(self, pai):
        self.Balanceamento = None
        self.bottomframe=Frame(pai)
        self.bottomframe.pack(side = BOTTOM)
        self.topframe=Frame(pai)
        self.topframe.pack(side=TOP)
        self.bottomtop=Frame(self.topframe)
        self.bottomtop.pack(side=BOTTOM)
        self.toptop=Frame(self.topframe)
        self.toptop.pack(side=TOP)
        
        self.text = Entry(self.toptop)
        self.text.bind("<Return>", self.Arvore)
        self.text.pack(side=LEFT)

        self.bDigite = Button(self.toptop)
        self.bDigite.bind("<Button-1>", self.Arvore)
        self.bDigite["text"] = "Valor"
        self.bDigite.pack(side=LEFT)
        
        self.bExibir = Button(self.bottomtop)
        self.bExibir["text"] = "Exibe Árvore AVL"
        self.bExibir.bind("<Button-1>", self.impArvore)
        self.bExibir.pack(side=LEFT)
        
        self.bBalanco = Button(self.bottomtop)
        self.bBalanco["text"] = "Balanceamento"
        self.bBalanco.bind("<Button-1>", self.impBal)
        self.bBalanco.pack(side=LEFT)
        
        self.c1 = Canvas(self.bottomframe, width=1000, height=550)
        self.c1.pack()

    def Arvore(self, *args):
        try:
            valor = int(self.text.get())
        except Exception:
            return
        print(valor)
        if self.Balanceamento == None:
            self.Balanceamento = Balanceamento()
            self.raiz = self.Balanceamento.criarNo(valor)
        else:
            self.Balanceamento.insValor(self.raiz, valor)
        self.grafArvore(False)

    def impArvore(self, *args):
        if self.Balanceamento != None:
            self.grafArvore(False)

    def impBal(self, *args):
        if self.Balanceamento != None:
            self.grafArvore(True)

    def grafArvore(self, comFB = False):
        self.HORIZONTAL = 1000
        self.VERTICAL = 600
        self.tamanho = 30
        self.c1.delete(ALL)
        self.c1.create_rectangle(
            0, 0, self.HORIZONTAL, self.VERTICAL, fill="black")
        self.xmax = self.c1.winfo_width() - 50
        self.ymax = self.c1.winfo_height()
        self.numero_linhas = self.Balanceamento.profundidadeMax(self.raiz)
        aux1 = int(self.xmax / 2 + 20)
        aux3 = int(self.ymax / (self.numero_linhas + 1))
        self.grafNo(self.raiz, aux1, aux3, aux1, aux3, 1, comFB)

    def grafNo(self, noh, posAX, posAY, posX, posY, linha, comFB = False):
        if noh == None:
            return
        numColunas = 2**(linha + 1)
        aux1 = int(posX - self.tamanho / 2)
        aux3 = int(posY - self.tamanho / 2)
        aux2 = int(posX + self.tamanho + 20 / 2)
        aux4 = int(posY + self.tamanho + 20/ 2)
        self.c1.create_oval(aux1, aux3, aux2, aux4, fill="gray")
        self.c1.create_line(posAX, posAY, posX, posY, fill="gray")
        

        rotulo = str(noh.dado)

        if comFB :
            rotulo = str(noh.balanco())
            
        self.c1.create_text(posX+10, posY+10, text=str(rotulo))
        posAX, posAY = posX, posY
        dx = self.xmax / numColunas
        dy = self.ymax / (self.numero_linhas + 1)
        posX = posAX + dx
        posY = posAY + dy
        self.grafNo(noh.direito, posAX, posAY, posX, posY, linha + 1, comFB)
        posX = posAX - dx
        self.grafNo(noh.esquerdo, posAX, posAY, posX, posY, linha + 1, comFB)

if __name__ == "__main__":
    root = Tk(None, None, "Árvore AVL")
    root.geometry("1000x600")
    ap = Grafico(root)
    root.mainloop()