basic-data-structure/AVLTree.py at master · haoliangyu/basic-data-structure · 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
from BinSearchTree import BinSearchTree
from BinNode import BinNode

class AVLTree(BinSearchTree):

    def insert(self, data):
        parent = { 'node': None }
        x = self.search(data, parent)
        if x is not None:
            return x

        x = BinNode(data, parent['node'])
        if x.data < parent['node'].data:
            parent['node'].left_child = x
        else:
            parent['node'].right_child = x

        self.set_size(self.size() + 1)

        node = parent['node']
        while node is not None:
            if not self.is_avl_balanced(node):
                is_left_child = None
                if not node.is_root():
                    is_left_child = node.is_left_child()

                rebalanced = self.__rotate_at(self.taller_child(self.taller_child(node)))

                if is_left_child is None:
                    self.root = rebalanced
                elif is_left_child:
                    rebalanced.parent.left_child = rebalanced
                else:
                    rebalanced.parent.right_child = rebalanced
                break
            else:
                self.update_height(node)

            node = node.parent

        return x

    def remove(self, data):
        parent = { 'node': None }
        x = self.search(data, parent)
        if x is None:
            return False

        self.remove_at(x, parent)
        self.set_size(self.size() - 1)

        node = parent['node']
        while node is not None:
            if not self.is_avl_balanced(node):
                is_left_child = None
                if not node.is_root():
                    is_left_child = node.is_left_child()

                rebalanced = self.__rotate_at(self.taller_child(self.taller_child(node)))

                if is_left_child is None:
                    self.root = rebalanced
                elif is_left_child:
                    rebalanced.parent.left_child = rebalanced
                else:
                    rebalanced.parent.right_child = rebalanced

            self.update_height(node)
            node = node.parent

        return True


    def is_balanced(self, node):
        return self.stature(node.left_child) == self.stature(node.right_child)

    def balance_factor(self, node):
        return self.stature(node.left_child) - self.stature(node.right_child)

    def is_avl_balanced(self, node):
        return -2 < self.balance_factor(node) < 2

    def stature(self, node):
        if node is None:
            return 0

        left_stature = self.stature(node.left_child)
        right_stature = self.stature(node.right_child)

        return (left_stature if left_stature > right_stature else right_stature) + 1

    def taller_child(self, node):
        stature_left = self.stature(node.left_child)
        stature_right = self.stature(node.right_child)

        if stature_left < stature_right:
            return node.right_child
        elif stature_left > stature_right:
            return node.left_child
        else:
            return node.left_child if node.is_left_child() else node.right_child

    def __connect34(self, a, b, c, t0, t1, t2, t3):
        a.left_child = t0
        if t0 is not None:
            t0.parent = a

        a.right_child = t1
        if t1 is not None:
            t1.parent = a

        self.update_height(a)

        c.left_child = t2
        if t2 is not None:
            t2.parent = c

        c.right_child = t3
        if t3 is not None:
            t3.parent = c

        self.update_height(c)

        b.left_child = a
        b.right_child = c

        a.parent = b
        c.parent = b

        self.update_height(b)

        return b

    def __rotate_at(self, v):
        p = v.parent
        g = p.parent

        if p.is_right_child():
            if v.is_right_child():
                p.parent = g.parent
                return self.__connect34(g, p, v, g.left_child, p.left_child, v.left_child, v.right_child)
            else:
                v.parent = g.parent
                return self.__connect34(g, v, p, g.left_child, v.left_child, v.right_child, p.right_child)
        else:
            if v.is_left_child():
                p.parent = g.parent
                return self.__connect34(v, p, g, v.left_child, v.right_child, p.right_child, g.right_child)
            else:
                v.parent = g.parent
                return self.__connect34(p, v, g, p.left_child, v.left_child, v.right_child, g.right_child)