AVL tree

import time
import random
import sys
sys.setrecursionlimit(3000)


def timing(f):
    def wrap(*args):
        time1 = time.time()
        ret = f(*args)
        time2 = time.time()
        print('%s took %0.3f s' % (f.__name__, (time2-time1)))
    return wrap


# I measure the height below a node!
def height(node):
    if node is None:
        return -1
    return max(height(node.left), height(node.right)) + 1


class Node:
    def __init__(self, value, parent=None, left=None, right=None):
        self.value = value
        self.parent = parent
        self.__left = left
        self.__right = right

    @property
    def left(self):
        return self.__left

    @left.setter
    def left(self, child):
        if child is not None:
            child.parent = self
        self.__left = child

    @property
    def right(self):
        return self.__right

    @right.setter
    def right(self, child):
        if child is not None:
            child.parent = self
        self.__right = child

    def add_children(self, value):

        if value < self.value:
            if self.left is None:
                self.left = Node(value, parent=self)
            else:
                self.left.add_children(value)

        else:
            if self.right is None:
                self.right = Node(value, parent=self)
            else:
                self.right.add_children(value)

    def get_height_scale(self):
        return height(self.left) - height(self.right)

    def print_height(self):
        return print('height', height(self.left), height(self.right))

    # All rotations starts from the highest node!!!

    def right_right_rotate(self):
        a = self

        b = a.left
        R = a.right

        L = b.left
        C = b.right

        if a.parent:
            highest_position_in_highest_parent = 'left' if a.value < a.parent.value else 'right'
            if highest_position_in_highest_parent == 'left':
                a.parent.left = b
            else:
                a.parent.right = b
        else:
            b.parent = None


        b.left = L
        b.right = a

        a.left = C
        a.right = R

    def left_left_rotate(self):
        a = self

        L = a.left
        b = a.right

        R = b.right
        C = b.left

        if a.parent:
            highest_position_in_highest_parent = 'left' if a.value < a.parent.value else 'right'
            if highest_position_in_highest_parent == 'left':
                a.parent.left = b
            else:
                a.parent.right = b

        else:
            b.parent = None

        b.left = a
        b.right = R

        a.left = L
        a.right = C

    # Don't forget to do LL or RR rotate after me on high-level bridge!
    def left_right_rotate(self):
        a = self

        R = a.right
        b = self.left

        L = b.left
        c = b.right
        M = c.left
        N = c.right

        if a.parent:
            a_position_in_a_parent = 'left' if a.value < a.parent.value else 'right'
            if a_position_in_a_parent == 'left':
                a.parent.left = c
            else:
                a.parent.right = c
        else:
            c.parent = None

        c.left = b
        c.right = a

        b.left = L
        b.right = M

        a.left = N
        a.right = R

    def right_left_rotate(self):
        a = self

        L = a.left
        b = a.right

        c = b.left
        R = b.right

        M = c.left
        N = c.right

        if a.parent:
            a_position_in_a_parent = 'left' if a.value < a.parent.value else 'right'
            if a_position_in_a_parent == 'left':
                a.parent.left = c
            else:
                a.parent.right = c
        else:
            c.parent = None

        c.left = a
        c.right = b

        a.left = L
        a.right = M

        b.left = N
        b.right = R


class AVLTreeKit:
    def __init__(self):
        self.top = None

    # I measure the height below a node!
    def height(self, node):
        if node is not None:
            return height(node)
        return 0

    def dynamic_find_top(self, node):
        if node.parent is None:
            self.top = node
        else:
            return self.dynamic_find_top(node.parent)

    def balance_parent(self, node):

        if height(node.right) - height(node.left) == 2 and height(node.right.left) <= height(node.right.right):
            node.left_left_rotate()
        elif height(node.right) - height(node.left) == 2 and height(node.right.left) > height(node.right.right):
            node.right_left_rotate()
        elif height(node.left) - height(node.right) == 2 and height(node.left.right) <= height(node.left.left):
            node.right_right_rotate()
        elif height(node.left) - height(node.right) == 2 and height(node.left.right) > height(node.left.left):
            node.left_right_rotate()

        if node.parent is not None:
            return self.balance_parent(node.parent)

    def uncheck_top_add_node(self, node, value):
        if value < node.value:
            if node.left is None:
                node.left = Node(value, parent=node)
                self.balance_parent(node)
            else:
                return self.uncheck_top_add_node(node.left, value)
        else:
            if node.right is None:
                node.right = Node(value, parent=node)
                self.balance_parent(node)
            else:
                return self.uncheck_top_add_node(node.right, value)

    def add_node(self, value):
        if self.top is None:
            self.top = Node(value)
        else:
            self.uncheck_top_add_node(self.top, value)
            self.dynamic_find_top(self.top)

    def dev_across(self, node_list, test=False):
        queue_list = []
        # [queue_list.extend([node.left, node.right]) for node in node_list if node is not None]
        for node in node_list:
            if node is not None and node != '<b>':
                if test:
                    if node.right is not None and node.right.value < node.value:
                        raise AssertionError('Error!,', 'Node:', node.value, 'Left:', node.left.value, 'Right:',
                                             node.right.value)
                    if node.left is not None and node.left.value > node.value:
                        raise AssertionError('Error!,', 'Node:', node.value, 'Left:', node.left.value, 'Right:',
                                             node.right.value)
                queue_list.extend([node.left, node.right])
            else:
                queue_list.extend(['<b>', '<b>'])
        for el in node_list:
            try:
                print(el.value, end=' ')
            except AttributeError:
                if el == '<b>':
                    print('<>', end=' ')
                else:
                    print('None', end=' ')
        print()
        if any([elem != '<b>' for elem in queue_list]):
            return self.dev_across(queue_list)

    def across(self, test=False):
        return self.dev_across([self.top], test)

    def info(self, test=False):
        self.across(test)
        self.top.print_height()


@timing
def main():
    kit = AVLTreeKit()
    for i in range(99, 0, -1):
        kit.add_node(random.randint(0, i))
    kit.info(test=True)

if __name__ == '__main__':
    main()

'''
    RR example
    
    kit = AVLTreeKit()
    kit.add_node(5)
    kit.add_node(3)
    kit.add_node(1)
    kit.across()
    print('scale', kit.top.get_height_scale())

    print('\nmaking RR rotate\n')

    kit.top.right_right_rotate()
    kit.dynamic_find_top(kit.top)
    kit.across()
    print('scale', kit.top.get_height_scale())
    
    LL example
    
    kit = AVLTreeKit()
    kit.add_node(5)
    kit.add_node(8)
    kit.add_node(9)
    kit.across()
    print('scale', kit.top.get_height_scale())

    print('\nmaking LL rotate\n')

    kit.top.left_left_rotate()
    kit.dynamic_find_top(kit.top)
    kit.across()
    print('scale', kit.top.get_height_scale())
    
    LR example
    
    kit = AVLTreeKit()
    kit.add_node(5)
    kit.add_node(3)
    kit.top.left.right = Node(4, parent=kit.top.left)
    kit.across()
    print('scale', kit.top.get_height_scale())

    print('\nmaking LR rotate\n')
    middle = kit.top.left.right
    kit.top.left_right_rotate()
    kit.top = middle
    kit.across()
    print('scale', kit.top.get_height_scale())
    
    RL example
    
    kit = AVLTreeKit()
    kit.add_node(5)
    kit.add_node(9)
    kit.top.right.left = Node(7, parent=kit.top.left)
    kit.across()
    print('scale', kit.top.get_height_scale())

    print('\nmaking RL rotate\n')

    kit.top.right_left_rotate()
    kit.dynamic_find_top(kit.top)
    kit.across()
    print('scale', kit.top.get_height_scale())
    
'''