DSA-AVL-Visualizer/bst.py at main · programmerlia/DSA-AVL-Visualizer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from node import Node
# ===========================================================
# AVL CLASS
# ===========================================================
class AVLTree:

    def insert(self, root, key):
        if root is None:
            return Node(key)

        #inserts
        if key < root.key:
            root.left = self.insert(root.left, key)
        elif key > root.key:
            root.right = self.insert(root.right, key)
        else:
            return root

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

        #Get balance
        balance = self.get_balance(root)


        # LL Case
        if balance > 1 and key < root.left.key:
            return self.ll_case(root)

        # RR Case
        if balance < -1 and key > root.right.key:
            return self.rr_case(root)

        # LR Case
        if balance > 1 and key > root.left.key:
            root.left = self.rr_case(root.left)
            return self.ll_case(root)

        # RL Case
        if balance < -1 and key < root.right.key:
            root.right = self.ll_case(root.right)
            return self.rr_case(root)

        return root

    def get_height(self,root):
        if not root:
            return 0
        return root.height

    def get_balance(self, root):
        if not root:
            return 0
        return self.get_height(root.left) - self.get_height(root.right)

    #rotations
    def ll_case(self, _1):
        _2 = _1.left
        child = _2.right
        _2.right = _1
        _1.left = child

        _2.height = 1+max(self.get_height(_2.left), self.get_height(_2.right))
        _1.height = 1+max(self.get_height(_1.left), self.get_height(_1.right))

        return _2

    def rr_case(self, _1):
        _2 = _1.right
        child = _2.left
        _2.left = _1
        _1.right = child

        _2.height = 1+max(self.get_height(_2.left), self.get_height(_2.right))
        _1.height = 1+max(self.get_height(_1.left), self.get_height(_1.right))
        return _2