Python实现自平衡二叉树AVL
# -*- 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)
很惭愧，只做了一点微小的工作
