def insert(self, root, key):
# Step 1 - Perform normal BST
root.left = self.insert(root.left, key)
root.right = self.insert(root.right, key)
# Step 2 - Update the height of the ancestor node
root.height = 1 + max(self.getHeight(root.left),
self.getHeight(root.right))
# Step 3 - Get the balance factor
balance = self.getBalance(root)
# Step 4 - If the node is unbalanced,
# then try out the 4 cases
if balance > 1 and key < root.left.val:
return self.rightRotate(root)
if balance < -1 and key > root.right.val:
return self.leftRotate(root)
if balance > 1 and key > root.right.val:
root.left = self.leftRotate(root.left)
return self.rightRotate(root)
if balance < -1 and key < root.right.val:
self.right = self.rightRotate(root.right)
return self.leftRotate(root)
z.height = 1 + max(self.getHeight(z.left),
y.height = 1 + max(self.getHeight(y.left),
def rightRotate(self, z):
z.height = 1 + max(self.getHeight(z.left),
y.height = 1 + max(self.getHeight(y.left),
def getHeight(self, root):
def getBalance(self, root):
return self.getHeight(root.left) - self.getHeight(root.right)