avl_kth/avl.py

class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.height = 1
        self.size = 1

    def __repr__(self) -> str:
        return f'{self.val}({self.size})'


class AVL_Tree(object):
    @staticmethod
    def insert(root, key):
        if not root:
            return TreeNode(key)
        elif key < root.val:
            root.left = AVL_Tree.insert(root.left, key)
        else:
            root.right = AVL_Tree.insert(root.right, key)

        root.height = 1 + max(AVL_Tree.getHeight(root.left),
                              AVL_Tree.getHeight(root.right))
        root.size = 1 + \
            AVL_Tree.getSize(root.left) + AVL_Tree.getSize(root.right)
        balance = AVL_Tree.getBalance(root)

        if balance > 1 and key <= root.left.val:
            return AVL_Tree.rightRotate(root)

        if balance < -1 and key >= root.right.val:
            return AVL_Tree.leftRotate(root)

        if balance > 1 and key >= root.left.val:
            root.left = AVL_Tree.leftRotate(root.left)
            return AVL_Tree.rightRotate(root)

        if balance < -1 and key <= root.right.val:
            root.right = AVL_Tree.rightRotate(root.right)
            return AVL_Tree.leftRotate(root)

        return root

    @staticmethod
    def leftRotate(z):
        y = z.right
        T2 = y.left

        y.left = z
        z.right = T2

        z.height = 1 + max(AVL_Tree.getHeight(z.left),
                           AVL_Tree.getHeight(z.right))
        y.height = 1 + max(AVL_Tree.getHeight(y.left),
                           AVL_Tree.getHeight(y.right))
        ty = AVL_Tree.getSize(y.right)
        tz = AVL_Tree.getSize(z.left)
        z.size -= (ty + 1)
        y.size += (tz + 1)
        return y

    @staticmethod
    def rightRotate(z):
        y = z.left
        T3 = y.right

        y.right = z
        z.left = T3

        z.height = 1 + max(AVL_Tree.getHeight(z.left),
                           AVL_Tree.getHeight(z.right))
        y.height = 1 + max(AVL_Tree.getHeight(y.left),
                           AVL_Tree.getHeight(y.right))

        ty = AVL_Tree.getSize(y.left)
        tz = AVL_Tree.getSize(z.right)
        z.size -= (ty + 1)
        y.size += (tz + 1)
        return y

    @staticmethod
    def getHeight(root):
        if not root:
            return 0
        return root.height

    @staticmethod
    def getSize(root):
        if not root:
            return 0
        return root.size

    @staticmethod
    def getBalance(root):
        if not root:
            return 0
        return AVL_Tree.getHeight(root.left) - AVL_Tree.getHeight(root.right)

    @staticmethod
    def preOrder(root, depth=0):
        if not root:
            return
        print("{0}{1}".format(' ' * depth, root))
        AVL_Tree.preOrder(root.left, depth+1)
        AVL_Tree.preOrder(root.right, depth+1)

    @staticmethod
    def getitem(root, idx):
        #print(root, idx)
        s = AVL_Tree.getSize(root.left)
        if s == idx:
            return root.val
        elif idx > s:
            return AVL_Tree.getitem(root.right, idx - s - 1)
        else:
            return AVL_Tree.getitem(root.left, idx)

    @staticmethod
    def index(root, val):
        if not root:
            return 0
        if root.val < val:
            return AVL_Tree.index(root.right, val) + AVL_Tree.getSize(root.left) + 1
        else:
            return AVL_Tree.index(root.left, val)
# Driver program to test above function

root = None


z = [6,21,-27,17,-20,3,1,-2,10,2,23,15,-3,1,9,19,-9,-24,-30,-26,-13,23,2,-10,20,0,27,24,-28,26,0,-29,-16,0,12,-28,7,1,22,-23,20,-22,-11,7,-10,-5,27,27,0,19,-9,28,-2,6,23,-9,-9,1,8,-15]

for i in z:
    root = AVL_Tree.insert(root, i)

AVL_Tree.preOrder(root)

for i in range(root.size):
    print(AVL_Tree.getitem(root, i))