AVL樹的python實現

AVL樹是帶有平衡條件的二叉查找樹，一般要求每個節點的左子樹和右子樹的高度最多差1(空樹的高度定義為-1)。

在高度為h的AVL樹中，最少的節點數S(h)由S(h)=S(h-1)+S(h-2)+1得出，其中S(0)=1，S(1)=2。

如上圖，分別為高度為0,1,2,3的AVL樹所需要的最少節點數。

1.AVL樹的實現，遍歷與查找操作與二叉查找樹相同。

class Node(object):
    def __init__(self,key):
        self.key=key
        self.left=None
        self.right=None
        self.height=0
class AVLTree(object):
    def __init__(self):
        self.root=None
    def find(self,key):
        if self.root is None:
            return None
        else:
            return self._find(key,self.root)
    def _find(self,key,node):
        if node is None:
            return None
        elif key<node.key:
            return self._find(key,self.left)
        elif key>node.key:
            return self._find(key,self.right)
        else:
            return node
    def findMin(self):
        if self.root is None:
            return None
        else:
            return self._findMin(self.root)
    def _findMin(self,node):
        if node.left:
            return self._findMin(node.left)
        else:
            return node
    def findMax(self):
        if self.root is None:
            return None
        else:
            return self._findMax(self.root)
    def _findMax(self,node):
        if node.right:
            return self._findMax(node.right)
        else:
            return node
    def height(self,node):
        if node is None:
            return -1
        else:
            return node.height

2.AVL樹的插入操作

插入一個節點可能會破壞AVL樹的平衡，可以通過旋轉操作來進行修正。

插入一個節點后，只有從插入節點到根節點的路徑上的節點的平衡可能被改變。我們需要找出第一個破壞了平衡條件的節點，稱之為K。K的兩顆子樹的高度差2。

不平衡有四種情況：

1.對K的左兒子的左子樹進行一次插入

2.對K的左兒子的右子樹進行一次插入

3.對K的右兒子的左子樹進行一次插入

4.對K的右兒子的右子樹進行一次插入

情況1與4是對稱的，需要進行一次單旋轉操作，清況2與3需要一次雙旋轉操作。

情況1：

 

    def singleLeftRotate(self,node):
        k1=node.left
        node.left=k1.right
        k1.right=node
        node.height=max(self.height(node.right),self.height(node.left))+1
        k1.height=max(self.height(k1.left),node.height)+1
        return k1

　　

情況4：

 

def singleRightRotate(self,node):
    k1=node.right
    node.right=k1.left
    k1.left=node
    node.height=max(self.height(node.right),self.height(node.left))+1
    k1.height=max(self.height(k1.right),node.height)+1
    return k1

情況3:

 

相當於進行了兩次單旋轉。

def doubleRightRotate(self,node):
    node.right=self.singleLeftRotate(node.right)
    return self.singleRightRotate(node)

情況2：

與情況3類似，都是進行了2次單旋轉。

def doubleLeftRotate(self,node):
    node.left=self.singleRightRotate(node.left)
    return self.singleLeftRotate(node)

一系列插入操作：

插入代碼如下：

    def put(self,key):
        if not self.root:
            self.root=Node(key)
        else:
            self.root=self._put(key,self.root)
    def _put(self,key,node):
        if node is None:
            node=Node(key)
        elif key<node.key:
            node.left=self._put(key,node.left)
            if (self.height(node.left)-self.height(node.right))==2:
                if key<node.left.key:
                    node=self.singleLeftRotate(node)
                else:
                    node=self.doubleLeftRotate(node)
            
        elif key>node.key:
            node.right=self._put(key,node.right)
            if (self.height(node.right)-self.height(node.left))==2:
                if key<node.right.key:
                    node=self.doubleRightRotate(node)
                else:
                    node=self.singleRightRotate(node)
        
        
        node.height=max(self.height(node.right),self.height(node.left))+1
        return node

3.AVL樹的刪除操作：

刪除操作比較復雜，如有錯誤，請指正。

1.當前節點為要刪除的節點且是樹葉（無子樹），直接刪除，當前節點（為None）的平衡不受影響。

2.當前節點為要刪除的節點且只有一個左兒子或右兒子，用左兒子或右兒子代替當前節點，當前節點的平衡不受影響。

3.當前節點為要刪除的節點且有左子樹右子樹:如果右子樹高度較高，則從右子樹選取最小節點，將其值賦予當前節點，然后刪除右子樹的最小節點。如果左子樹高度較高，則從左子樹選取最大節點，將其值賦予當前節點，然后刪除左子樹的最大節點。這樣操作當前節點的平衡不會被破壞。

4.當前節點不是要刪除的節點，則對其左子樹或者右子樹進行遞歸操作。當前節點的平衡條件可能會被破壞，需要進行平衡操作。

如上圖，25為當前節點，左子樹刪除17后平衡條件被破壞，需要根據當前節點(25)的右子樹(30)的左子樹(28)高度是否高於右子樹(35)的高度進行判斷，若高於，進行雙旋轉，否則進行單旋轉

    def delete(self,key):
        self.root=self.remove(key,self.root)
    def remove(self,key,node):
        if node is None:
            raise KeyError,'Error,key not in tree'
        elif key<node.key:
            node.left=self.remove(key,node.left)
            if (self.height(node.right)-self.height(node.left))==2:
                if self.height(node.right.right)>=self.height(node.right.left):
                    node=self.singleRightRotate(node)
                else:
                    node=self.doubleRightRotate(node)
            node.height=max(self.height(node.left),self.height(node.right))+1
            
                
        elif key>node.key:
            node.right=self.remove(key,node.right)
            if (self.height(node.left)-self.height(node.right))==2:
                if self.height(node.left.left)>=self.height(node.left.right):
                    node=self.singleLeftRotate(node)
                else:
                    node=self.doubleLeftRotate(node)
            node.height=max(self.height(node.left),self.height(node.right))+1
        
        elif node.left and node.right:
            if node.left.height<=node.right.height:
                minNode=self._findMin(node.right)
                node.key=minNode.key
                node.right=self.remove(node.key,node.right)
            else:
                maxNode=self._findMax(node.left)
                node.key=maxNode.key
                node.left=self.remove(node.key,node.left)
            node.height=max(self.height(node.left),self.height(node.right))+1
        else:
            if node.right:
                node=node.right
            else:
                node=node.left
        
        return node

　　

全部代碼：

class Node(object):
    def __init__(self,key):
        self.key=key
        self.left=None
        self.right=None
        self.height=0
class AVLTree(object):
    def __init__(self):
        self.root=None
    def find(self,key):
        if self.root is None:
            return None
        else:
            return self._find(key,self.root)
    def _find(self,key,node):
        if node is None:
            return None
        elif key<node.key:
            return self._find(key,self.left)
        elif key>node.key:
            return self._find(key,self.right)
        else:
            return node
    def findMin(self):
        if self.root is None:
            return None
        else:
            return self._findMin(self.root)
    def _findMin(self,node):
        if node.left:
            return self._findMin(node.left)
        else:
            return node
    def findMax(self):
        if self.root is None:
            return None
        else:
            return self._findMax(self.root)
    def _findMax(self,node):
        if node.right:
            return self._findMax(node.right)
        else:
            return node
    def height(self,node):
        if node is None:
            return -1
        else:
            return node.height
    
    def singleLeftRotate(self,node):
        k1=node.left
        node.left=k1.right
        k1.right=node
        node.height=max(self.height(node.right),self.height(node.left))+1
        k1.height=max(self.height(k1.left),node.height)+1
        return k1
    def singleRightRotate(self,node):
        k1=node.right
        node.right=k1.left
        k1.left=node
        node.height=max(self.height(node.right),self.height(node.left))+1
        k1.height=max(self.height(k1.right),node.height)+1
        return k1
    def doubleLeftRotate(self,node):
        node.left=self.singleRightRotate(node.left)
        return self.singleLeftRotate(node)
    def doubleRightRotate(self,node):
        node.right=self.singleLeftRotate(node.right)
        return self.singleRightRotate(node)
    def put(self,key):
        if not self.root:
            self.root=Node(key)
        else:
            self.root=self._put(key,self.root)
    def _put(self,key,node):
        if node is None:
            node=Node(key)
        elif key<node.key:
            node.left=self._put(key,node.left)
            if (self.height(node.left)-self.height(node.right))==2:
                if key<node.left.key:
                    node=self.singleLeftRotate(node)
                else:
                    node=self.doubleLeftRotate(node)
            
        elif key>node.key:
            node.right=self._put(key,node.right)
            if (self.height(node.right)-self.height(node.left))==2:
                if key<node.right.key:
                    node=self.doubleRightRotate(node)
                else:
                    node=self.singleRightRotate(node)
        
        
        node.height=max(self.height(node.right),self.height(node.left))+1
        return node
        
    def delete(self,key):
        self.root=self.remove(key,self.root)
    def remove(self,key,node):
        if node is None:
            raise KeyError,'Error,key not in tree'
        elif key<node.key:
            node.left=self.remove(key,node.left)
            if (self.height(node.right)-self.height(node.left))==2:
                if self.height(node.right.right)>=self.height(node.right.left):
                    node=self.singleRightRotate(node)
                else:
                    node=self.doubleRightRotate(node)
            node.height=max(self.height(node.left),self.height(node.right))+1
            
                
        elif key>node.key:
            node.right=self.remove(key,node.right)
            if (self.height(node.left)-self.height(node.right))==2:
                if self.height(node.left.left)>=self.height(node.left.right):
                    node=self.singleLeftRotate(node)
                else:
                    node=self.doubleLeftRotate(node)
            node.height=max(self.height(node.left),self.height(node.right))+1
        
        elif node.left and node.right:
            if node.left.height<=node.right.height:
                minNode=self._findMin(node.right)
                node.key=minNode.key
                node.right=self.remove(node.key,node.right)
            else:
                maxNode=self._findMax(node.left)
                node.key=maxNode.key
                node.left=self.remove(node.key,node.left)
            node.height=max(self.height(node.left),self.height(node.right))+1
        else:
            if node.right:
                node=node.right
            else:
                node=node.left
        
        return node