import sys
#创建树节点
class
TreeNode(object):
def __init__(self,key):
self.key=key
self.left=None
self.right=None
self.height=1
class
AVLTree(object):
#插入节点
def insert_node(self,root,key):
#找到位置并插入节点
if
not root:
return
TreeNode(key)
elif key<root.key:
root.left=self.insert_node(root.left,key)
else
:
root.right=self.insert_node(root.right,key)
root.height=1+max(self.getHeight(root.left),
self.getHeight(root.right))
#更新节点
balanceFactor=self.getBalance(root)
if
balanceFactor>1:
if
key<root.left.key:
return
self.rightRotate(root)
else
:
root.left=self.leftRotate(root.left)
return
self.rightRotate(root)
if
balanceFactor<-1:
if
key>root.right.key:
return
self.leftRotate(root)
else
:
root.right=self.rightRotate(root.right)
return
self.leftRotate(root)
return
root
#删除节点
def delete_node(self,root,key):
#找到要删除的节点并删除
if
not root:
return
root
elif key<root.key:
root.left=self.delete_node(root.left,key)
elif key>root.key:
root.right=self.delete_node(root.right,key)
else
:
if
root.left is None:
temp=root.right
root=None
return
temp
elif root.right is None:
temp=root.left
root=None
return
temp
temp=self.getMinValueNode(root.right)
root.key=temp.key
root.right=self.delete_node(root.right,
temp.key)
if
root is None:
return
root
#更新节点
root.height=1+max(self.getHeight(root.left),
self.getHeight(root.right))
balanceFactor=self.getBalance(root)
#平衡树
if
balanceFactor>1:
if
self.getBalance(root.left)>=0:
return
self.rightRotate(root)
else
:
root.left=self.leftRotate(root.left)
return
self.rightRotate(root)
if
balanceFactor<-1:
if
self.getBalance(root.right)<=0:
return
self.leftRotate(root)
else
:
root.right=self.rightRotate(root.right)
return
self.leftRotate(root)
return
root
#左旋转
def leftRotate(self,z):
y=z.right
T2=y.left
y.left=z
z.right=T2
z.height=1+max(self.getHeight(z.left),
self.getHeight(z.right))
y.height=1+max(self.getHeight(y.left),
self.getHeight(y.right))
return
y
#右旋转
def rightRotate(self,z):
y=z.left
T3=y.right
y.right=z
z.left=T3
z.height=1+max(self.getHeight(z.left),
self.getHeight(z.right))
y.height=1+max(self.getHeight(y.left),
self.getHeight(y.right))
return
y
#获取节点高度
def getHeight(self,root):
if
not root:
return
0
return
root.height
#平衡节点
def getBalance(self,root):
if
not root:
return
0
return
self.getHeight(root.left)-self.getHeight(root.right)
def getMinValueNode(self,root):
if
root is None
or
root.left is None:
return
root
return
self.getMinValueNode(root.left)
def preOrder(self,root):
if
not root:
return
print
(
"{0}"
.format(root.key),
end
=
""
)
self.preOrder(root.left)
self.preOrder(root.right)
#输出avl树
def printHelper(self,currPtr,indent,last):
if
currPtr!=None:
sys.stdout.write(indent)
if
last:
sys.stdout.write(
"R----"
)
indent+=
""
else
:
sys.stdout.write(
"L----"
)
indent+=
"|"
print
(currPtr.key)
self.printHelper(currPtr.left,indent,False)
self.printHelper(currPtr.right,indent,True)
myTree=AVLTree()
root=None
nums=[33,13,52,9,21,61,8,11]
for
num in nums:
root=myTree.insert_node(root,num)
myTree.printHelper(root,
""
,True)
key=13
root=myTree.delete_node(root,key)
print
(
"After Deletion:"
)
myTree.printHelper(root,
""
,True)