4
\$\begingroup\$

I have tried implementing an AVL Tree on my own, based on visualising it. But i'm unsure how many testcases it work with, and how efficient it is. Are there any ways to make it efficient, and compact?

class Node:
    def __init__(self, value, parent=None):
        """
        Every Node has a value, and a height based on how far it is from the root node
        """
        self.value = value
        self.parent = parent
        self.left = None
        self.right = None

class AVLTree:
    def __init__(self, List=None):
        """
        initialising the root which is the first node which is None in an empty Tree

        If List keyword argument is provided, then it will insert every value in the List.
        """
        self.root = None
        if List:
            for item in List:
                self.insert(item)

    def inorder(self, current_node):
        """
        Recursive Inorder Traversal of the Binary Tree, to display from smallest to largest, the values in the tree
        Ex:
            1 2 3 5 6 9
        """
        if current_node:  
            self.inorder(current_node.left)
            print(current_node.value, end=' ')
            self.inorder(current_node.right)

    def get_depth(self, node):
        """
        Finds the depth of a specific node, which means the levels of nodes beneath it, also recursive.
        Ex:                               height:
                    O                       0
                /        \                  
               O          O                 1
            /     \    /      \
           O       O  O        O            2
                              / 
                             O              3

            The Depth of root node (height 0) = 3, since there are 3 levels of nodes beneath it
            The Depth of left node at (height 1) = 1, since there is only 1 level of nodes beneath it
            The Depth of right node at (height 1) = 2, since there are ``  2 ``   ``   ``     ``    ``
            The Depth of left node of (height 2) = 0, since there are no nodes beneath it
            The Depth of a Non Existent Node = 0
        """
        # Null node has 0 depth.
        if node is None:
            return 0

        # Get the depth of the left and right subtree 
        # using recursion.
        leftDepth = self.get_depth(node.left)
        rightDepth = self.get_depth(node.right)

        # Choose the larger one and add the root to it.
        if leftDepth > rightDepth:
            return leftDepth + 1
        elif leftDepth < rightDepth:
            return rightDepth + 1
        else: # A Node with no children has a depth of 0
            return 0

    def get_balance_factor(self, node):
        """
        Balance_factor is the difference between the heights of the two child node of the node

        Rebalancing occurs when the balance_factor is greater than a difference of 1 (or) less than a difference of -1

        usage: balance_factor = get_balance_factor(node)
        which is left_child_height - right_child_height

        If balanced, balance_factor = -1/0/1:
                O           O          O         O      
               / \         / \        / \       / \ 
              O   O       O   O      O   O     O   O
                         / \        /               \ 
                        O   O      O                 O   

        If balance_factor > 1:
                O    <--- Node
               / \ 
              O   O   Height = 2
             / \ 
            O   O     Height = 3
           / \ 
          O   O       Height = 4  
            Difference of L and R = 2 (4-2)

        If balance_factor < 1:
                O   <--- Node
               / \ 
              O   O        Height = 2
                 / \        
                O   O      Height = 3
                   / \ 
                  O   O    Height = 4
            Difference of L and R = -2 (2-4)


        """
        if node is None:
            return 
        return self.get_depth(node.left) - self.get_depth(node.right)


    def rotate_right(self, y):
        """
        Rotates right
        Ex:

              y                               x
             / \     Right Rotation          /  \
            x   T3   - - - - - - - >        T1   y 
           / \                                  / \
          T1  T2                              T2  T3

        changed:
        1) the left child of y became T2
        2) the right child of x became y
        3) y's parent now points to x
        4) T2's parent is now y if it exists
        5) x's parent is now y's parent

        """
        # Rotation
        parent = y.parent

        x = y.left
        T2 = x.right

        x.right = y
        y.left = T2

        if T2:
            T2.parent = y
        x.parent = parent

        # If there is no parent, then that means y was the root node
        if parent is None:
            self.root = x
        elif parent.left == y:
            parent.left = x
        elif parent.right == y:
            parent.right = x

    def rotate_left(self, x):
        """
        Rotates left
        Ex:

              y                               x
             / \                             /  \
            x   T3                          T1   y 
           / \       < - - - - - - -            / \
          T1  T2     Left Rotation            T2  T3

        changed:
        1) the left child of y became x
        2) the right child of x became T2
        3) x's parent now points to y
        4) T2's parent is now x if it exists
        5) y's parent is now x's parent

        """
        # Rotation
        parent = x.parent
        y = x.right
        T2 = y.left

        y.left = x
        x.right = T2

        if T2:
            T2.parent = x
        y.parent = parent

        # If there is no parent, then that means y was the root node
        if parent is None:
            self.root = y 
        elif parent.left == x:
            parent.left = y
        elif y.parent.right == x:
            parent.right = y

    def recursive_balancing_for_insertion(self, current_node, value):
        """
        As the name says, rebalances, by using comparisons between the value to be inserted and child node's value
        It slowly checks the balance factor of every node from the parent of the recently inserted node,
        after a node has been checked, it goes the parent of that node, until it reaches the root node, which doesn't have a parent /= None.

        Depending on the balance factor and comparison, it may perform a left rotation, or right rotation, or even both.
        Rotations shift the position of the nodes so that o(log N) operations will still be performed, and prevents skewed trees

        Cases:

        1) Left Left         2) Right Right
                O                  O
               /                    \ 
              O                      O
             /                        \ 
            O                          O

        2) Right Left        4) Left Right
                O                   O
               /                     \ 
              O                       O
               \                      /
                O                    O


        """
        if current_node is None:
            return

        # Getting the balance factor, which shows if the tree is balanced or not after insertion, by the difference of the heights between the left and right children of the node
        balance_factor = self.get_balance_factor(current_node)

        # Checks if the node is unbalanced, and if it is, perform one of the 4 cases
        # Case 1: Left Left
        if balance_factor > 1 and current_node.left:
            if value < current_node.left.value:
                self.rotate_right(current_node)

        # Case 2: Right Right
        if balance_factor < -1 and current_node.right:
            if value > current_node.right.value:
                self.rotate_left(current_node)

        # Case 3: Left Right
        if balance_factor > 1 and current_node.left:
            if value > current_node.left.value:
                self.rotate_left(current_node.left)
                self.rotate_right(current_node)

        # Case 4: Right Left
        if balance_factor < -1 and current_node.right:
            if value < current_node.right.value:
                self.rotate_right(current_node.right)
                self.rotate_left(current_node)

        self.recursive_balancing_for_insertion(current_node.parent, value)

    def recursive_balancing_for_deletion(self, current_node):
        """
        As the name says, rebalances, by using balance factor of it's child nodes
        It slowly checks the balance factor of every node from the parent of the recently inserted node,
        after a node has been checked, it goes the parent of that node, until it reaches the root node, which doesn't have a parent /= None.

        Depending on the balance factor and comparison, it may perform a left rotation, or right rotation, or even both.
        Rotations shift the position of the nodes so that o(log N) operations will still be performed, and prevents skewed trees

        Cases:

        1) Left Left         2) Right Right
                O                  O
               /                    \ 
              O                      O
             /                        \ 
            O                          O

        2) Right Left        4) Left Right
                O                   O
               /                     \ 
              O                       O
               \                      /
                O                    O

        """
        if current_node is None:
            return

        balance_factor = self.get_balance_factor(current_node)

        # Case 1 - Left Left
        if balance_factor > 1 and self.get_balance_factor(current_node.left) >= 0: 
            self.rotate_right(current_node) 

        # Case 2 - Right Right 
        if balance_factor < -1 and self.get_balance_factor(current_node.right) <= 0: 
            self.rotate_left(current_node) 

        # Case 3 - Left Right 
        if balance_factor > 1 and self.get_balance_factor(current_node.left) < 0: 
            self.rotate_left(current_node.left) 
            self.rotate_right(current_node) 

        # Case 4 - Right Left 
        if balance_factor < -1 and self.get_balance_factor(current_node.right) > 0: 
            self.rotate_right(current_node.right) 
            self.rotate_left(current_node)

        self.recursive_balancing_for_deletion(current_node.parent)

    def insert(self, value):
        """
        Inserts a value into the tree, by traversing through the tree till it finds a None.

        1) No balancing is done when inserting the root node.
        2) Balancing is done for any other node.
        """
        NewNode = Node(value)
        current_node = self.root

        # Traversing through the Tree till we find the right place to put the value
        if current_node is None:
            self.root = NewNode
        else:
            while True:
                if value < current_node.value:
                    NewNode.parent = current_node
                    #Left
                    if not current_node.left:
                        current_node.left = NewNode
                        break
                    NewNode.parent = current_node
                    current_node = current_node.left
                else:
                    NewNode.parent = current_node
                    #Right
                    if not current_node.right:
                        current_node.right = NewNode
                        break
                    NewNode.parent = current_node
                    current_node = current_node.right

            # Recursive balancing
            self.recursive_balancing_for_insertion(current_node, value)

    def lookup(self, value):
        """
        Searches for a value in the tree,
        if the value is found, returns True
        if the value isn't found, returns False

        """
        current_node = self.root

        if current_node is None:
            return None
        while current_node:
            current_value = current_node.value
            if value < current_value:
                current_node = current_node.left
            elif value > current_value:
                current_node = current_node.right
            else:
                return True
        return False

    def getMin(self, current_node):
        """
        returns the smallest value of a node in the current_node, by going left, until there is no more left child nodes

        Used to find the in order successor for a node, incase it has two children when removing.

        """
        while current_node.left is not None:
            current_node = current_node.left
        return current_node.value

    def remove(self, value, current_node=None):
        """
        Removes a value from the tree
        returns True, if value was deleted
        returns False, if value wasn't found/deleted

        Different position changes depending on 3 cases:
        1) Node has no children
            easiest to remove, just remove the parent nodes pointer to it
        2) Node has 1 child
            just point the next child of parent node to the next child of the node to be deleted instead
        3) Node has 2 children
            hardest to remove, first the in order successor of the node is found, by getting the smallest value in the right child node.
            Then the node value is replaced with the in order successor, and the inorder successor is recursively deleted.
        """
        if not current_node:
            current_node = self.root

        if not current_node: # if the root is None
            return None

        parent_node = None
        while current_node:
            # value of current node
            current_value = current_node.value

            # If value to be searched is less than the current node value, then move left
            if value < current_value:   
                parent_node = current_node
                current_node = current_node.left

            # If value to be searched is greater than the current node value, then move right
            elif value > current_value:
                parent_node = current_node
                current_node = current_node.right

            # if value to be searched is the current node's value
            else:
                # No Child
                if not current_node.left and not current_node.right:

                    # If there is no parent node, then it is the root node
                    if parent_node is None:
                        self.root = None
                    elif current_node == parent_node.left:
                        parent_node.left = None
                    else:
                        parent_node.right = None
                    self.recursive_balancing_for_deletion(current_node)

                # One Child
                elif current_node.right and not current_node.left:
                    if parent_node is None:
                        self.root = current_node.right
                    elif current_node == parent_node.left:
                        parent_node.left = current_node.right
                    else:
                        parent_node.right = current_node.right
                    current_node.right.parent = parent_node
                    current_node.right = None
                    self.recursive_balancing_for_deletion(current_node)

                elif current_node.left and not current_node.right:
                    if parent_node is None:
                        self.root = current_node.left
                    elif current_node == parent_node.left:
                        parent_node.left = current_node.left
                    else:
                        parent_node.right = current_node.left
                    current_node.left.parent = parent_node
                    current_node.left = None
                    self.recursive_balancing_for_deletion(current_node)

                # Two Child
                else:
                    # gets the successor of the current_node
                    in_order_successor = self.getMin(current_node.right)

                    # Removes the successor to replace the node to be removed
                    self.remove(in_order_successor, current_node)

                    #if the node to be removed is the root node
                    if parent_node is None:
                        self.root.value = in_order_successor
                    elif current_node == parent_node.left:
                        parent_node.left.value = in_order_successor
                    else:
                        parent_node.right.value = in_order_successor 

                    self.recursive_balancing_for_deletion(parent_node)

                current_node.parent = None
                return True # if removed

        return False # if value doesnt exist



test = AVLTree((1, 2, 3, 4, 5))        
#test.remove(3)
#test.remove(1)
#test.remove(2)
#test.remove(4)
#test.remove(5)
#print(test.lookup(5))
test.inorder(test.root)

\$\endgroup\$
0
\$\begingroup\$

I believe it's usual, and more efficient to store the balance of each node ( -1, 0, +1 ) in the node rather than computing it each time.

Another point : I think you could use some more elif statements, only one kind of rotation will apply.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.