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
from tree.tree import TreeNode


class AvlTree(object):
    """
    An avl tree.
    """

    def __init__(self):
        # Root node of the tree.
        self.node = None
        self.height = -1
        self.balance = 0

    def insert(self, key):
        """
        Insert new key into node
        """
        # Create new node
        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):
        """
        Re balance tree. After inserting or deleting a node,
        """
        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):
        """
        Update tree height
        """
        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):
        """
        Calculate tree balance factor

        """
        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):
        """
        Right rotation
        """
        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):
        """
        Left rotation
        """
        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):
        """
        In-order traversal of the tree
        """
        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
