imports treenodes from tree tree import treenode class avltreeobject def initself root node of the tree self node none self height 1 self balance 0 def insertself key create new node node treenodekey if not self node self node node self node left avltree self node right avltree elif key self node val self node left insertkey elif key self node val self node right insertkey self rebalance def rebalanceself self updateheightsrecursivefalse self updatebalancesfalse while self balance 1 or self balance 1 if self balance 1 if self node left balance 0 self node left rotateleft self updateheights self updatebalances self rotateright self updateheights self updatebalances if self balance 1 if self node right balance 0 self node right rotateright self updateheights self updatebalances self rotateleft self updateheights self updatebalances def updateheightsself recursivetrue if self node if recursive if self node left self node left updateheights if self node right self node right updateheights self height 1 maxself node left height self node right height else self height 1 def updatebalancesself recursivetrue if self node if recursive if self node left self node left updatebalances if self node right self node right updatebalances self balance self node left height self node right height else self balance 0 def rotaterightself newroot self node left node newleftsub newroot right node oldroot self node self node newroot oldroot left node newleftsub newroot right node oldroot def rotateleftself newroot self node right node newleftsub newroot left node oldroot self node self node newroot oldroot right node newleftsub newroot left node oldroot def inordertraverseself result if not self node return result result extendself node left inordertraverse result appendself node key result extendself node right inordertraverse return result an avl tree root node of the tree insert new key into node create new node re balance tree after inserting or deleting a node update tree height calculate tree balance factor right rotation left rotation in order traversal of the tree
from tree.tree import TreeNode


class AvlTree(object):
    

    def __init__(self):
        
        self.node = None
        self.height = -1
        self.balance = 0

    def insert(self, key):
        
        
        node = TreeNode(key)
        if not self.node:
            self.node = node
            self.node.left = AvlTree()
            self.node.right = AvlTree()
        elif key < self.node.val:
            self.node.left.insert(key)
        elif key > self.node.val:
            self.node.right.insert(key)
        self.re_balance()

    def re_balance(self):
        
        self.update_heights(recursive=False)
        self.update_balances(False)

        while self.balance < -1 or self.balance > 1:
            if self.balance > 1:
                if self.node.left.balance < 0:
                    self.node.left.rotate_left()
                    self.update_heights()
                    self.update_balances()
                self.rotate_right()
                self.update_heights()
                self.update_balances()

            if self.balance < -1:
                if self.node.right.balance > 0:
                    self.node.right.rotate_right()
                    self.update_heights()
                    self.update_balances()
                self.rotate_left()
                self.update_heights()
                self.update_balances()

    def update_heights(self, recursive=True):
        
        if self.node:
            if recursive:
                if self.node.left:
                    self.node.left.update_heights()
                if self.node.right:
                    self.node.right.update_heights()

            self.height = 1 + max(self.node.left.height,
                                  self.node.right.height)
        else:
            self.height = -1

    def update_balances(self, recursive=True):
        
        if self.node:
            if recursive:
                if self.node.left:
                    self.node.left.update_balances()
                if self.node.right:
                    self.node.right.update_balances()

            self.balance = self.node.left.height - self.node.right.height
        else:
            self.balance = 0

    def rotate_right(self):
        
        new_root = self.node.left.node
        new_left_sub = new_root.right.node
        old_root = self.node

        self.node = new_root
        old_root.left.node = new_left_sub
        new_root.right.node = old_root

    def rotate_left(self):
        
        new_root = self.node.right.node
        new_left_sub = new_root.left.node
        old_root = self.node

        self.node = new_root
        old_root.right.node = new_left_sub
        new_root.left.node = old_root

    def in_order_traverse(self):
        
        result = []

        if not self.node:
            return result

        result.extend(self.node.left.in_order_traverse())
        result.append(self.node.key)
        result.extend(self.node.right.in_order_traverse())
        return result