avl_tree.py

# AVl Quelle:https://www.programiz.com/dsa/avl-tree

# AVL tree implementation in Python


import sys


# Create a tree node
class TreeNode(object):
    def __init__(self, key, vertice, vertice_neighbour):
        self.key = key
        self.vertice = vertice
        self.vertice_neighbour = vertice_neighbour
        self.left = None
        self.right = None
        self.height = 1


class AVLTree(object):

    def find_edges(self, root, key, vertice_top):
        if not root:
            return vertice_top
        elif key < root.key:
            vertice_top = root
            return self.find_edges(root.left, key, vertice_top)
        elif key == root.key:
            return root
        else:
            if vertice_top.key < key:  # falls wir irgendwann im Buam nach links gegangen sind ist dieser Node der Nachfolger, ansonsten gibt es kein Nachfolger und die Ecke berührt den Container rand
                vertice_top = TreeNode(key, (0, key), (0, 0))
            return self.find_edges(root.right, key, vertice_top)
        # return ("test",vertice_top)

    # Function to insert a node
    def insert_node(self, root, key, vertice, vertice_neighbour):

        # Find the correct location and insert the node
        if not root:
            return TreeNode(key, vertice, vertice_neighbour)
        elif key < root.key:
            root.left = self.insert_node(root.left, key, vertice, vertice_neighbour)
        else:
            root.right = self.insert_node(root.right, key, vertice, vertice_neighbour)

        root.height = 1 + max(self.getHeight(root.left),
                              self.getHeight(root.right))

        # Update the balance factor and balance the tree
        balanceFactor = self.getBalance(root)
        if balanceFactor > 1:
            if key < root.left.key:
                return self.rightRotate(root)
            else:
                root.left = self.leftRotate(root.left)
                return self.rightRotate(root)

        if balanceFactor < -1:
            if key > root.right.key:
                return self.leftRotate(root)
            else:
                root.right = self.rightRotate(root.right)
                return self.leftRotate(root)

        return root

    # Function to delete a node
    def delete_node(self, root, key):

        # Find the node to be deleted and remove it
        if not root:
            return root
        elif key < root.key:
            root.left = self.delete_node(root.left, key)
        elif key > root.key:
            root.right = self.delete_node(root.right, key)
        else:
            if root.left is None:
                temp = root.right
                root = None
                return temp
            elif root.right is None:
                temp = root.left
                root = None
                return temp
            temp = self.getMinValueNode(root.right)
            root.key = temp.key
            root.vertice = temp.vertice
            root.vertice_neighbour = temp.vertice_neighbour
            root.right = self.delete_node(root.right,
                                          temp.key)
        if root is None:
            return root

        # Update the balance factor of nodes
        root.height = 1 + max(self.getHeight(root.left),
                              self.getHeight(root.right))

        balanceFactor = self.getBalance(root)

        # Balance the tree
        if balanceFactor > 1:
            if self.getBalance(root.left) >= 0:
                return self.rightRotate(root)
            else:
                root.left = self.leftRotate(root.left)
                return self.rightRotate(root)
        if balanceFactor < -1:
            if self.getBalance(root.right) <= 0:
                return self.leftRotate(root)
            else:
                root.right = self.rightRotate(root.right)
                return self.leftRotate(root)
        return root

    # Function to perform left rotation
    def leftRotate(self, z):
        y = z.right
        T2 = y.left
        y.left = z
        z.right = T2
        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 y

    # Function to perform right rotation
    def rightRotate(self, z):
        y = z.left
        T3 = y.right
        y.right = z
        z.left = T3
        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 y

    # Get the height of the node
    def getHeight(self, root):
        if not root:
            return 0
        return root.height

    # Get balance factore of the node
    def getBalance(self, root):
        if not root:
            return 0
        return self.getHeight(root.left) - self.getHeight(root.right)

    def getMinValueNode(self, root):
        if root is None or root.left is None:
            return root
        return self.getMinValueNode(root.left)

    def preOrder(self, root):
        if not root:
            return
        print("{0} ".format(root.key), end="")
        self.preOrder(root.left)
        self.preOrder(root.right)

    def preOrder_array(self, root, array, key):
        if not root:
            return
            # print("{0} ".format(root.key), end="")
        if root.key <= key:
            array.append((root.vertice))
        self.preOrder_array(root.left, array, key)
        self.preOrder_array(root.right, array, key)
        return array

    # Print the tree
    def printHelper(self, currPtr, indent, last):
        if currPtr != None:
            sys.stdout.write(indent)
            if last:
                sys.stdout.write("R----")
                indent += "     "
            else:
                sys.stdout.write("L----")
                indent += "|    "
            print(currPtr.key)
            self.printHelper(currPtr.left, indent, False)
            self.printHelper(currPtr.right, indent, True)