【二叉樹】用python實現AVL樹

AVL樹是平衡的二叉搜索樹,其任意節點的平衡因子(左子樹高度-右子樹高度)始終爲-1、0或1。AVL樹的結構能夠解決普通BTS由於有序插入導致的退化爲鏈表的情況。

AVL樹在二叉搜索樹實現的基礎上,主要在添加和刪除節點時更新子樹相關節點的平衡因子,同時對不平衡的子樹進行旋轉操作。其機制較爲複雜,本文先介紹相關的背景知識,在做代碼的實現。

1. 背景知識
1.1 節點平衡因子的調整順序

與普通二叉搜索樹不同的是:添加或刪除節點後相關節點的平衡因子需做出調整。

  • 添加操作

其基本規則包括:
(1)若該節點爲左節點,則其父節點的平衡因子加1;若該節點爲右節點,則其父節點的平衡因子減1;然後繼續向上回溯父節點;
(2)若父節點的平衡因子調整爲0,則說明以該父節點爲根節點的子樹達到了平衡,但高度不變,因此後續的祖先節點無需在做更新;
(3)若父節點的平衡因子仍爲-1或1,則繼續重複(1);
(4)若父節點的平衡因子爲-2或2,則以該父節點爲根節點的子樹進行旋轉操作(具體細節見1.2),並調整旋轉節點的平衡因子(具體細節見1.3)。旋轉後的子樹相較於原子樹,高度不變,因此後續的祖先節點無需做更新;
(5)若有必要,一直回溯到AVL樹的根節點爲止。

顯然,可以通過遞歸實現上述調整,其條件終止條件包括:
條件一:到達根節點;
條件二:某節點的平衡因子達到0;
條件三:某節點的平衡因子達到-2或2。

  • 刪除操作

刪除操作是在原AVL樹基礎上任意刪除某個節點的操作。根據普通二叉搜索樹中刪除操作的解讀,無論刪除的是根、中間還是葉子節點,其最終都可以轉換爲葉子節處的刪除。

根據原AVL樹的定義,若刪除葉子節點所在最底層的子樹原來是完全平衡的,則其父節點在刪除後的平衡因子爲1或-1,此時整棵子樹的高度不變,因此無需繼續上溯。

而若刪除葉子節點所在最底層的子樹原來並非是完全平衡的(即原父節點的平衡因子爲1或-1),則刪除後的平衡因子可能變爲0、-2或2。若爲0,則說明原子樹父節點的平衡因子爲1或-1;刪除後子樹的高度變化,因此需繼續回溯。而對於-2和2,需要先採用左旋或右旋操作將子樹平衡,然後視新的父節點的平衡因子選擇繼續回溯或者停止。

基於上面的分析,可採用類似於增加節點的遞歸操作實現調整,其條件終止條件包括:
條件一:到達根節點;
條件二:某節點的平衡因子達到-1或1。

1.2 子樹的旋轉

子樹包括兩種基本旋轉方式:左旋和右旋。

  • 左旋

當右子樹重(即根節點平衡因子爲-2)時,需要進行左旋。其具體具體操作包括:
(1)根節點右孩子節點變爲新的根節點,原根節點成爲新根節點的左孩子節點;
(2)若原根節點本身就具有左孩子節點,則該節點成爲原根節點的右孩子節點。
在這裏插入圖片描述

  • 右旋

當左子樹重(即根節點平衡因子爲2)時,需要進行右旋。其具體具體操作包括:
(1)根節點左孩子節點變爲新的根節點,原根節點成爲新根節點的右孩子節點;
(2)若原根節點本身就具有右孩子節點,則該節點成爲原根節點的左孩子節點。
在這裏插入圖片描述

在刪除節點後的再平衡過程中,由於原樹已經處於平衡狀態,葉子節點的刪除(即使是非葉子節點,也會通過與後繼節點的交換轉換爲葉子節點的刪除)通過單次的左旋或右旋操作即可完成當前子樹的再平衡。

而在添加節點後的再平衡過程中,情況會略顯複雜,其可能涉及到雙旋轉的操作。具體而言,根據子樹情況可分爲四種變換:
情況1: 根節點平衡因子小於0,右孩子節點的平衡因子也小於0,直接在對根節點所在的樹進行左旋。
在這裏插入圖片描述
情況2:根節點平衡因子小於0,而右孩子節點的平衡因子大於0,先對右子樹進行右旋,然後對根節點所在的樹進行左旋。
在這裏插入圖片描述
情況3: 根節點平衡因子大於0,左孩子節點的平衡因子也大於0,直接在對根節點所在的樹進行右旋。
在這裏插入圖片描述
情況4: 根節點平衡因子大於0,而左孩子節點的平衡因子小於0,先對左子樹進行左旋,然後對根節點所在的樹進行右旋。
在這裏插入圖片描述

1.3 旋轉後相關節點平衡因子的更新公式

值得注意的是:在每步的旋轉過程中,只有發生位置變換的根節點和成爲新的根節點的那個原孩子節點的平衡因子需要進行更新。而以其他節點爲根節點的子樹結構並未發生變動,因此節點上的平衡因子也無需變化。那麼具體如何更新呢?我們可以推導下左旋和右旋下的更新公式。

  • 左旋
    在這裏插入圖片描述

上圖爲一個典型的左旋示意圖。我們不妨令旋轉前,根節點(對應點B)的左、右子樹高度分別爲H1H_1H2H_2;右孩子節點D(左旋後成爲新根節點)爲根節點的子樹的左、右子樹高度分別爲h1h_1h2h_2,則節點B和和節點D的平衡因子分別爲:rB=H1H2rD=h1h2r_B=H_1-H_2\\ r_D=h_1-h_2且滿足H2=max(h1,h2)+1H_2=\max(h_1,h_2)+1

旋轉後:點B的左子樹不變,而右子樹變成了點D的左子數,因此,其平衡因子變爲:rˉB=H1h1=(H1H2)+(H2h1)=rB+max(0,h2h1)+1=rBmin(0,h1h2)+1=rBmin(0,rD)+1\begin{aligned}\bar r_B&=H_1-h_1 \\&=(H_1-H_2)+(H_2-h_1)\\&=r_B+\max(0,h_2-h_1)+1 \\&=r_B-\min(0, h_1-h_2)+1\\&=r_B-\min(0,r_D)+1\end{aligned}
旋轉後:點D的右子樹不變,而左子樹變爲以B爲根節點的子樹,因此,其平衡因子變爲:rˉD=max(H1,h1)+1h2=max(H1,h1)h1+1+h1h2=max(H1h1,0)+1+(h1h2)=max(rˉB,0)+1+rD\begin{aligned}\bar r_D&=\max(H_1,h_1)+1-h_2\\&=\max(H_1,h_1)-h_1+1+h_1-h_2\\&=\max(H_1-h_1,0)+1+(h_1-h_2)\\&=\max(\bar r_B,0)+1+r_D\end{aligned}

上面兩式中的rBr_BrˉB\bar r_B可泛化爲原根節點,左旋後的左孩子節點;而rDr_DrˉD\bar r_D可泛化爲原右孩子節點和左旋轉後的新的根節點。

  • 右旋
    在這裏插入圖片描述

與左旋類似,則節點E和節點D的平衡因子分別爲:rE=H1H2rC=h1h2r_E=H_1-H_2\\ r_C=h_1-h_2且滿足H1=max(h1,h2)+1H_1=\max(h_1,h_2)+1

旋轉後:點E的右子樹不變,而左子樹變成了點C的右子樹,因此,其平衡因子變爲:rˉE=h2H2=(h2H1)+(H1H2)=rEmax(rC,0)1\begin{aligned}\bar r_E&=h_2-H_2 \\&=(h_2-H_1)+(H_1-H_2)\\&=r_E-\max(r_C,0)-1\end{aligned}

旋轉後:點C的左子樹不變,而右子樹變爲以E爲根節點的子樹,因此,其平衡因子變爲:rˉC=h1(max(h2,H2)+1)=h1h2+h2+min(h2,H2)1=(h1h2)+min(0,h2H2)1=rC+min(0,rˉE)1\begin{aligned}\bar r_C&=h_1-(\max(h_2,H_2)+1)\\&=h_1-h_2+h_2+\min(-h_2,-H_2)-1\\&=(h_1-h_2)+\min(0, h_2-H_2)-1\\&=r_C+\min(0, \bar r_E)-1\end{aligned}
上面兩式中的rEr_ErˉE\bar r_E可泛化爲原根節點,左旋後的右孩子節點;而rCr_CrˉC\bar r_C可泛化爲原左孩子節點和右旋轉後的新的根節點。

2. 代碼實現

AVL樹可繼承普通二叉搜索樹實現,通過設置新的節點屬性:平衡因子;同時覆寫添加和刪除節點時相應節點平衡因子的更新以及相關子樹的旋轉操作即可。

2.1 增加平衡因子屬性
class TreeNode(object):
    def __init__(self, key, value, left=None, right=None, parent=None):
        self.key = key
        self.value = value
        self.leftChild = left
        self.rightChild =right
        self.parent = parent
        self.balanceFactor = 0   #  新增屬性,用於記錄節點的平衡因子,初始值爲0
   
    # pass
2.2 繼承BST構造AVL Tree
class AVLTree(BinarySearchTree):
    def __init__(self):
        super(AVLTree, self).__init__()
2.3 定義左旋和右旋轉操作
class AVLTree(BinarySearchTree):
	"""
	左旋和右旋轉操作和BTS一樣,這裏僅做說明
	"""
    def leftRotate(self, node):
       """
       左旋操作:(1)右孩子節點爲新根節點,原根節點爲左孩子節點;(2)若原右孩子節點本存在左孩子節點,則成爲新左孩子節點的右孩子節點
       注意指針調整順序、平衡因子的更新以及邊界條件
       """
       newRoot = node.rightChild
       node.rightChild = newRoot.leftChild
       if newRoot.leftChild:
           newRoot.leftChild.parent = node

       if node.isLeftChild():
           newRoot.parent = node.parent
           node.parent.leftChild = newRoot
       elif node.isRightChild():
           newRoot.parent = node.parent
           node.parent.rightChild = newRoot
       else:
           newRoot.parent = None
           self.root = newRoot

       newRoot.leftChild = node
       node.parent = newRoot

       node.balanceFactor = node.balanceFactor - min(0, newRoot.balanceFactor) + 1
       newRoot.balanceFactor = newRoot.balanceFactor + max(node.balanceFactor, 0) + 1

   def rightRotate(self, node):
       """
       右旋操作:(1)左孩子節點爲新根節點,原根節點爲右孩子節點;(2)若原左孩子節點本存在右孩子節點,則成爲新右孩子節點的左孩子節點
       注意指針調整順序、平衡因子的更新以及邊界條件
       """
       newRoot = node.leftChild
       node.leftChild = newRoot.rightChild
       if newRoot.rightChild:
           newRoot.rightChild.parent = node

       if node.isLeftChild():
           newRoot.parent = node.parent
           node.parent.leftChild = newRoot
       elif node.isRightChild():
           newRoot.parent = node.parent
           node.parent.rightChild = newRoot
       else:
           newRoot.parent = None
           self.root = newRoot

       newRoot.rightChild = node
       node.parent = newRoot

       node.balanceFactor = node.balanceFactor - max(newRoot.balanceFactor, 0) - 1
       newRoot.balanceFactor = newRoot.balanceFactor + min(0, node.balanceFactor) - 1
2.4 覆寫添加元素相關方法
class AVLTree(BinarySearchTree):
	# pass
	
    def _put(self, node: TreeNode, key, value):
        """
        覆寫添加節點操作,再原有基礎上,更新平衡因子
        """
        if key == node.key:
            node.value = value
        elif key < node.key:
            if not node.hasLeftChild():
                node.leftChild = TreeNode(key, value, parent=node)
                self.put_updateBalanceFactor(node.leftChild)
            else:
                node = node.leftChild
                self._put(node, key, value)
        else:
            if not node.hasRightChild():
                node.rightChild = TreeNode(key, value, parent=node)
                self.put_updateBalanceFactor(node.rightChild)
            else:
                node = node.rightChild
                self._put(node, key, value)

    def put_updateBalanceFactor(self, node):
        """
        遞歸更新平衡因子並視情況進行子樹旋轉,其原則包括:
        (1)由當前節點向上回溯其父節點,視當前節點爲左/右節點更新父節點的平衡因子;
        (2)終止條件包括:1)某祖先節點平衡因子調整爲0;2)達到根節點;3)某祖先節點平衡因子調整爲-2或2,進行旋轉操作
        """
        if node.balanceFactor > 1 or node.balanceFactor < -1:
            self.put_rotateTree(node)
            return

        if not node.isRoot():
            if node.isLeftChild():
                node.parent.balanceFactor += 1

            else:
                node.parent.balanceFactor -= 1

            if node.parent.balanceFactor != 0:
                self.put_updateBalanceFactor(node.parent)

    def put_rotateTree(self, node):
        """
        以某節點爲根的子樹的旋轉操作,視該節點的平衡因子以及對應左孩子或右孩子節點的平衡因子選擇對應的旋轉方式
        因爲觸犯條件爲平衡因子2或者-2所以在對應子樹上必有孫子節點

        """
        if node.balanceFactor > 1:
            if node.leftChild.balanceFactor >= 0:    # 情況1:直接右旋
                self.rightRotate(node)

            else:                                   # 情況2:先左旋右子樹,再右旋當前樹
                self.leftRotate(node.rightChild)
                self.rightRotate(node)

        elif node.balanceFactor < -1:
            if node.rightChild.balanceFactor <= 0:    # 情況3: 直接左旋
                self.leftRotate(node)

            else:
                self.rightRotate(node.leftChild)      # 情況4: 先右旋左子樹,再左旋當前樹
                self.leftRotate(node)
   # pass 
2.5 覆寫刪除元素相關方法
class AVLTree(BinarySearchTree):
	"""
	其和BST的區別在於刪除節點後,要對節點的父/祖先節點的平衡因子進行更新及子樹的旋轉
	注意刪除操作和添加中平衡因子更新的方式和終止條件的區別
	"""
	# pass
    def delete(self, key):
        """
        AVL樹的delete在BST基礎上新增了平衡因子操作
        """
        if not self.root:  # 空樹
            raise KeyError("It is empty bts!")
        if self.size == 1:  # 只有根節點
            if self.root.key == key:
                self.root = None
                self.size = 0
            else:
                raise KeyError("No key in the bts!")

        else:
            node = self._get(self.root, key)
            if not node:  # 未找到key
                raise KeyError("No key in the bts!")

            elif node.isLeaf():  # 葉子節點,且肯定不爲根節點,因此肯定有父節點
                if node.isLeftChild():
                    node.parent.leftChild = None
                    self.delete_updateBalanceFactor(node)

                else:
                    node.parent.rightChild = None
                    self.delete_updateBalanceFactor(node)

            elif not node.hasBothChildren():  # 只有一個子樹
                if node.hasLeftChild():  # 情況1:只有左子樹,且子樹只有1個根節點
                    if node.isRoot():  # 情況1.1: 情況當前節點爲根節點
                        node.replaceNodeData(node.leftChild.key, node.leftChild.value, node.leftChild.leftChild,
                                             node.leftChild.rightChild)
                    elif node.isLeftChild():  # 情況1.2: 當前節點爲左節點
                        node.parent.leftChild = node.leftChild
                        node.leftChild.parent = node.parent
                        self.delete_updateBalanceFactor(node)

                    else:  # 情況1.3: 當前節點右節點
                        node.parent.rightChild = node.leftChild
                        node.leftChild.parent = node.parent
                        self.delete_updateBalanceFactor(node)

                else:  # 情況2:只有右子樹
                    if node.isRoot():  # 情況2.1: 情況當前節點爲根節點
                        node.replaceNodeData(node.rightChild.key, node.rightChild.value, node.rightChild.leftChild,
                                             node.rightChild.rightChild)
                    elif node.isLeftChild():  # 情況1.2: 當前節點爲左節點
                        node.parent.leftChild = node.rightChild
                        node.rightChild.parent = node.parent
                        self.delete_updateBalanceFactor(node)

                    else:  # 情況1.3 當前節點爲右節點
                        node.parent.rightChild = node.rightChild
                        node.rightChild.parent = node.parent
                        self.delete_updateBalanceFactor(node)

            else:  # 存在左右兩棵子樹,所以其後繼節點必然爲右子樹上的最小值
                successor = self._findSuccessor(node)  # 該後繼節點最多隻有一棵子樹
                if successor.isLeaf():  # 葉子節點
                    if successor.isLeftChild():
                        successor.parent.leftChild = None
                    else:
                        successor.parent.rightChild = None

                else:  # 左子樹的最低層左節點,或者右子樹的根節點
                    successor.parent.rightChild = successor.rightChild
                    successor.rightChild.parent = successor.parent

                node.replaceNodeData(successor.key, successor.value, node.leftChild, node.rightChild)
                self.delete_updateBalanceFactor(successor)   # 物理本質上真正刪除的是後繼節點

            self.size -= 1

    def delete_updateBalanceFactor(self, node):

        if not node.isRoot():
            if node.isLeftChild():
                node.parent.balanceFactor -= 1

            else:
                node.parent.balanceFactor += 1

            if node.parent.balanceFactor == 0:
                self.delete_updateBalanceFactor(node.parent)

            if node.parent.balanceFactor < -1 or node.parent.balanceFactor > 1:
                self.delete_rotateTree(node.parent)
                self.delete_updateBalanceFactor(node.parent.parent)    # 經過旋轉,當前節點的父節點變爲了新的孩子節點

        if node.balanceFactor in [1, -1]:       # 主要順序,要先更新父節點平衡因子
            return

    def delete_rotateTree(self, node):
        """
        以某節點爲根的子樹的旋轉操作,視該節點的平衡因子選擇對應的旋轉方式
        因爲BST刪除節點而來,所以只有兩種基本的旋轉方式
        """
        if node.balanceFactor > 1:
            self.rightRotate(node)

        elif node.balanceFactor < -1:
            self.leftRotate(node)
   # pass 
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章