如何在Python中实现avl树运算

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)