Python实现自平衡二叉树AVL

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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
# -*- coding: utf-8 -*-
from enum import Enum
 
#参考http://blog.csdn.net/niteip/article/details/11840691/
#参考https://www.cnblogs.com/suimeng/p/4560056.html
#todo 还没有考虑高度的增减,只考虑了平衡因子
#todo 加上非递归遍历二叉树
 
class BFStatus(Enum):
    # 左子树的高度减去右子树的高度 RH,EH,LH分别表示右子树较高,左右子树等高,左子树较高
    LH = 1
    EH = 0
    RH = -1
 
class TreeNode(object):
    def __init__(self):
        self.lson = None
        self.rson = None
        self.pnode = None
        self.data = None
        self.freq = 0
        self.hgt = -1
        self.bfstatus = BFStatus().EH
 
class AVLTree(object):
    def __init__(self,root_node):
        self.root = TreeNode
 
    def height(self,node):
        if node is not None:
            return node.hgt
        return -1
 
    def max(self,campa,campb):
        if campa > campb:
            return campa
        else:
            return campb
 
    def single_rotate_right(self, node):  # 右单旋转 RR
        parent_node = node.pnode
        node1 = node.lson
        node.lson = node1.rson
        node1.rson = node
        if parent_node.lson == node:
            parent_node.lson = node1
        else:
            parent_node.rson = node1
        node1.pnode = parent_node
        node.pnode = node1
 
    def single_rotate_left(self, node):  # 左单旋转 LL
        parent_node = node.pnode
        node1 = node.rson
        node.rson = node1.lson
        node1.lson = node
        if parent_node.lson == node:
            parent_node.lson = node1
        else:
            parent_node.rson = node1
        node1.pnode = parent_node
        node.pnode = node1
 
    def LeftBalance(self,node):
        lchild = node.lson
        BFStatus_code = BFStatus()
        if lchild.bfstatus == BFStatus_code.EH:
            node.bfstatus = lchild.bfstatus = BFStatus_code.LH
            lchild.bfstatus = BFStatus.RH
            self.single_rotate_right(node)
        elif lchild.bfstatus == BFStatus_code.LH:#LL型
            node.bfstatus = lchild.bfstatus = BFStatus.EH
            self.single_rotate_right(node)
        elif lchild.bfstatus == BFStatus.RH:#LR型
            rchild = lchild.rson
            if rchild.bfstatus == BFStatus.EH:
                node.bfstatus = BFStatus.RH
                lchild.bfstatus = BFStatus.EH
            elif rchild.bfstatus == BFStatus.LH:
                node.bfstatus = BFStatus.EH
                lchild.bfstatus = BFStatus.LH
            elif rchild.bfstatus == BFStatus.RH:
                node.bfstatus = BFStatus.EH
                lchild.bfstatus = BFStatus.EH
            rchild.bfstatus = BFStatus.EH
            self.single_rotate_left(lchild)
            self.single_rotate_right(node)
 
    def RightBalance(self,node):
        rchild = node.rson
        BFStatus_code = BFStatus()
        if rchild.bfstatus == BFStatus_code.RH:
            node.bfstatus = node.bfstatus = BFStatus_code.EH
            self.single_rotate_left(node)
        elif rchild.bfstatus == BFStatus_code.EH:#LL型
            node.bfstatus = rchild.bfstatus = BFStatus.RH
            rchild.bfstatus = BFStatus.LH
            self.single_rotate_left(node)
        elif rchild.bfstatus == BFStatus.LH:#LR型
            lchild = rchild.lson
            if lchild.bfstatus == BFStatus.LH:
                node.bfstatus = BFStatus.EH
                rchild.bfstatus = BFStatus.RH
            elif lchild.bfstatus == BFStatus.RH:
                node.bfstatus = BFStatus.LH
                rchild.bfstatus = BFStatus.EH
            elif lchild.bfstatus == BFStatus.EH:
                node.bfstatus = BFStatus.EH
                rchild.bfstatus = BFStatus.EH
            lchild.bfstatus = BFStatus.EH
            self.single_rotate_right(rchild)
            self.single_rotate_left(node)
 
 
    def insertpri(self,node,data,stack_list,taller_param=True):
        if node == None:
            node = TreeNode()
            node.data = data
            node.pnode = stack_list[-1]
        else:
            stack_list.append(node)
            if node.data < data:
                (sign,taller) = self.insertpri(node.lson,data,stack_list)
                if sign == False:
                    return (False,taller_param)
 
                if taller == True:
 
                    if node.bfstatus == BFStatus().LH:
                        self.LeftBalance(node)
                        taller_param = False
                    elif node.bfstatus == BFStatus().EH:
                        node.bfstatus =  BFStatus().LH
                        taller_param = True
                    elif node.bfstatus == BFStatus().RH:
                        node.bfstatus =  BFStatus().EH
                        taller_param = False
 
            elif node.data > data:
                stack_list.append(node)
                if node.data < data:
                    (sign, taller) = self.insertpri(node.rson, data, stack_list)
                    if sign == False:
                        return (False, taller_param)
 
                    if taller == True:
                        if node.bfstatus == BFStatus().LH:
                            node.bfstatus = BFStatus().EH
                            taller_param = False
                        elif node.bfstatus == BFStatus().EH:
                            node.bfstatus = BFStatus().RH
                            taller_param = True
                        elif node.bfstatus == BFStatus().RH:
                            self.RightBalance(node)
                            taller_param = False
            else:
                node.freq += 1
 
        return (True,taller_param)
 
    def insert(self,data):
        stack_list = []
        self.insertpri(self.root,data,stack_list)
 
 
    def searchpri(self,node,data):
        if node is None:
            return None
 
        elif node.data > data:
            self.searchpri(node.lson,data)
 
        elif node.data < data:
            self.searchpri(node.rson,data)
 
        else:
            return node
 
    def search(self,data):
        self.searchpri(self.root,data)
 
    def go_far_left(self,node):
        if node.lson is None:
            return node
        else:
            self.go_far_left(node.lson)
 
    def go_far_right(self,node):
        if node.rson is None:
            return node
        else:
            self.go_far_right(node.rson)
 
    def delete_left(self,node,bfchild):
        #当bf为-1或1变为0,或者孩子为空时说明子树高降低
        if node.lson is None or (bfchild != BFStatus().EH and node.lson.bfstatus == BFStatus().EH):
            #左子树树高降低
            if node.bfstatus == BFStatus().EH:
                node.bfstatus = BFStatus().RH
            elif node.bfstatus == BFStatus().LH:
                node.bfstatus = BFStatus().EH
            elif node.bfstatus == BFStatus().RH:#原本右子树比较高
                self.RightBalance(node)
 
    def delete_right(self,node,bfchild):
        #当bf为LH或RH变为EH,或者孩子为空时说明子树高降低
        if node.rson is None or (bfchild != BFStatus().EH and node.rson.bfstatus == BFStatus().EH):
            #左子树树高降低
            if node.bfstatus == BFStatus().EH:
                node.bfstatus = BFStatus().LH
            elif node.bfstatus == BFStatus().RH:
                node.bfstatus = BFStatus().EH
            elif node.bfstatus == BFStatus().LH:#原本左子树比较高
                self.LeftBalance(node)
 
    def deletepri(self,node,data):
        bfchild = BFStatus().EH
        if node == None:
            return None
        if node.data > data:
            bfchild = node.lson.bfstatus
            node.lson = self.deletepri(node.lson,data)
            self.delete_left(node,bfchild)
 
        elif node.data < data:
            bfchild = node.rson.bfstatus
            node.rson = self.deletepri(node.rson, data)
            self.delete_left(node, bfchild)
 
        else:
            if node.lson is not None:   #不是叶子结点并且具有直接前驱
                far_right_node = self.go_far_right(node.lson)#直接前驱
                node.data = far_right_node.data
                node.lson = self.deletepri(node.lson,far_right_node.data)#可以确定,删除的节点为当前节点的左子树的某一个节点
                self.delete_left(node,bfchild)
            elif node.rson is not None:
                far_left_node = self.go_far_left(node.rson)
                node.data = far_left_node.data
                node.rson = self.deletepri(node.rson,far_left_node.data)
                self.delete_right(node,bfchild)
            else:#叶子结点
                node = None
 
        return node
 
    def delete(self,data):
        self.deletepri(self.root,data)
 
 
    def insubtree(self,node):
        if node is None:
            return
        self.insubtree(node.lson)
        print(node.data)
        self.insubtree(node.rson)
 
    def traversal(self):
        self.insubtree(self.root)
 
if __name__ == '__main__':
    root_node = TreeNode()
    tree_obj = AVLTree(root_node)

  

posted @   巫谢  阅读(1546)  评论(0编辑  收藏  举报
点击右上角即可分享
微信分享提示