PAT 甲级 1066 Root of AVL Tree

1066 Root of AVL Tree

模拟平衡二叉树的过程，输出最后生成的树的根节点

class Node:
    def __init__(self,data,height=1,left=None,right=None):
        self.data=data
        self.left=left
        self.right=right
        self.height=height

    def getHeight(self):
        if self==None:
            return 0
        else:
            return max(self.left.getHeight() if self.left!=None else 0,self.right.getHeight() if self.right!=None else 0)+1

    def updateHeight(self):    #更新树高
        self.height=max(self.left.getHeight() if self.left!=None else 0,self.right.getHeight() if self.right!=None else 0)+1

    def getBalanceFactor(self):
        return (self.left.getHeight() if self.left!=None else 0)-(self.right.getHeight() if self.right!=None else 0)

    def leftRotation(self): #左旋
        tmp=self.right
        self.right=tmp.left if tmp.left!=None else None
        tmp.left=self
        tmp.updateHeight()
        self.updateHeight()
        return tmp

    def rightRotation(self):    #右旋
        tmp=self.left
        self.left=tmp.right if tmp.right!=None else None
        tmp.right=self
        tmp.updateHeight()
        self.updateHeight()
        return tmp

    def insert(self,x): #插入操作
        if x.data<self.data:
            if self.left==None:
                self.left=x
            else:
                self.left=self.left.insert(x)
            self.updateHeight() #更新树高
            if self.getBalanceFactor()==2:  #在左子树中插入节点使得整棵树不平衡
                if self.left.getBalanceFactor()==1: #左子树的左子树中插入节点使得整棵树不平衡
                    self=self.rightRotation()
                elif self.left.getBalanceFactor()==-1:  #左子树的右子树中插入节点使得整棵树不平衡
                    self.left=self.left.leftRotation()
                    self=self.rightRotation()
        elif x.data>self.data:
            if self.right==None:
                self.right=x
            else:
                self.right=self.right.insert(x)
            self.updateHeight()
            if self.getBalanceFactor()==-2: #在右子树中插入节点使得整棵树不平衡
                if self.right.getBalanceFactor()==-1:   #在右子树的右子树中插入节点使得整棵树不平衡
                    self=self.leftRotation()
                elif self.right.getBalanceFactor()==1:    #在右子树的左子树中插入节点使得整棵树不平衡
                    self.right=self.right.rightRotation()
                    self=self.leftRotation()
        return self

n=int(input())
l=list(map(int,input().split()))
root=Node(l[0])
for i in range(1,n):
    root=root.insert(Node(l[i]))
print(root.data)