DEV Community

loading...
Cover image for Insertion in AVL tree

Insertion in AVL tree

Aya Bouchiha
Full stack web developer
Updated on ・3 min read

Hi, I'm Aya Bouchiha and today I'm going to talk about the insertion in AVL tree, but if you're not familiar with AVL trees, check these posts below :

Before we get started, I want to mention that Balance Factor
of a balanced node should be always {-1,0,1}

BalanceFactor=height(left sub-tree)-height(right sub-tree)

Insertion implementation in python from geeksforgeeks

"""
    this insertion implementation is from geeksforgeeks
    https://www.geeksforgeeks.org/avl-tree-set-1-insertion/
    (This code is contributed by Ajitesh Pathak)
"""
class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.height = 1

# AVL tree class which supports the
# Insert operation
class AVL_Tree(object):

    # Recursive function to insert key in
    # subtree rooted with node and returns
    # new root of subtree.
    def insert(self, root, key):

        # Step 1 - Perform normal BST
        if not root:
            return TreeNode(key)
        elif key < root.val:
            root.left = self.insert(root.left, key)
        else:
            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
        # Case 1 - Left Left
        if balance > 1 and key < root.left.val:
            return self.rightRotate(root)

        # Case 2 - Right Right
        if balance < -1 and key > root.right.val:
            return self.leftRotate(root)

        # Case 3 - Left Right
        if balance > 1 and key > root.left.val:
            root.left = self.leftRotate(root.left)
            return self.rightRotate(root)

        # Case 4 - Right Left
        if balance < -1 and key < root.right.val:
            root.right = self.rightRotate(root.right)
            return self.leftRotate(root)

        return root

    def leftRotate(self, z):

        y = z.right
        T2 = y.left

        # Perform rotation
        y.left = z
        z.right = T2

        # Update heights
        z.height = 1 + max(self.getHeight(z.left),
                         self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                         self.getHeight(y.right))

        # Return the new root
        return y

    def rightRotate(self, z):

        y = z.left
        T3 = y.right

        # Perform rotation
        y.right = z
        z.left = T3

        # Update heights
        z.height = 1 + max(self.getHeight(z.left),
                        self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                        self.getHeight(y.right))

        # Return the new root
        return y

    def getHeight(self, root):
        if not root:
            return 0

        return root.height

    def getBalance(self, root):
        if not root:
            return 0

        return self.getHeight(root.left) - self.getHeight(root.right)


    def preOrder(self, root):
        if not root:
            return
        print(root.val, end="\t")
        self.preOrder(root.left)
        self.preOrder(root.right)



# Driver program to test above function
myTree = AVL_Tree()
root = None
root = myTree.insert(root, 10)
root = myTree.insert(root, 20)
root = myTree.insert(root, 30)
root = myTree.insert(root, 40)
root = myTree.insert(root, 50)
root = myTree.insert(root, 25)
"""
     30
    /  \
  20   40
 /  \     \
10  25    50
"""
# Preorder Traversal
print("Preorder traversal of the","constructed AVL tree is")
myTree.preOrder(root)
Enter fullscreen mode Exit fullscreen mode

for more details check this article

References and useful resources

#day_20

Happy coding!

Discussion (0)