手擼二叉樹——AVL平衡二叉樹

牛初九發表於2024-10-17
二叉樹

還記得上一篇中我們遺留的問題嗎？我們再簡要回顧一下，現在有一顆空的二叉查詢樹，我們分別插入1，2，3，4，5，五個節點，那麼得到的樹是什麼樣子呢？這個不難想象，二叉樹如下：

樹的高度是4，並且資料結構上和連結串列沒有區別，查詢效能也和連結串列一致。如果我們將樹的結構改變一下呢？比如改成下面的樹結構，

那麼樹的高度變成了3，並且它也是一棵二叉查詢樹，樹的高度越低，查詢效能就越高，這是我們理想中的資料結構。如果想要樹的高度儘可能的低，那麼左右子樹的高度差就不能相差太多。這就引出了我們今天的主題AVL平衡二叉樹，AVL平衡二叉樹的定義為任意節點的左右子樹的高度差不能超過1。這樣就可以保證我們的這棵樹的高度保持在一個最低的狀態，這樣我們的查詢效能也是最優的。那麼我們如何在樹的變化時（也就是增加節點或刪除節點時），保證AVL平衡二叉樹的性質呢？下面我們就針對每一種情況進行分析。

左左單旋轉

我們先看看下面的例子，以下每一個例子都是最複雜的情況，完全覆蓋簡單的情況，所以我們把最複雜情況用程式碼實現了，那麼簡單的情況也會涵蓋在內。看下圖

上圖中，原本以k1為根節點的樹是一個AVL平衡二叉樹，這時，我們向樹中插入節點2，根據二叉查詢樹的性質，最後節點2插入的位置如上圖。插入節點後，我們每個節點分析一下，看看節點是否還符合AVL平衡二叉樹的性質。我們先看看節點3，插入節點2後，節點3的左子樹的高度是0，因為只有一個節點2。再看節點3的右子樹，右子樹為空，那麼高度為-1，這裡我們統一規定，如果節點為空，那麼高度為-1。節點3的左右子樹高度為1，符合AVL平衡二叉樹的性質，同理我們再看節點k2，左子樹高度為1，右子樹高度為0，高度差為1，也符合AVL平衡二叉樹。再看節點k1，左子樹k2的高度為2，右子樹的高度為0，相差為2，所以在節點k1處不滿足AVL平衡二叉樹的性質，我們要進行調整，使得以k1為根節點的樹變為一個AVL平衡二叉樹，我們要怎麼做呢？

由於左子樹的高度比較高，所以我們要將樹旋轉一下，用k2作根節點，k1作為k2的右子節點，旋轉後如圖所示：

旋轉後，以k2為根節點的新樹，是一棵AVL平衡二叉樹。這裡我們要特別注意一下節點5的位置，它的原始位置是k2的右子樹，而k2又是k1的左子樹，根據二叉查詢樹的性質，k2的右子樹中的值是大於k2，小於k1的。旋轉後，k2變成了根節點，k1變成k2的右子樹，那麼原k2的右子樹（節點5），變為k1的左子樹。那麼這棵樹根據二叉查詢樹的性質，還是大於k2，小於k1的，沒有變動，這是符合我們的預期的。透過上述的旋轉，我們得到的新樹是一棵AVL平衡二叉樹。

我們總結一下重要的點，為編碼做準備：

  1. 發現k1的左子樹比右子樹高度大於1；
  2. 發現k1的左子樹k2的左子樹高度大於k2的右子樹高度，這種稱作左-左情形。要做左側單旋轉。
  3. 將k2作為新樹的節點，k2的右子樹改為k1，k1的左子樹改為k2的右子樹。
  4. 更新k1和k2的高度。

完成上面的操作，我們得到一個新的AVL平衡二叉樹。下面我們進入具體編碼。

/**
 * 二叉樹節點
 * @param <T>
 */
public class BinaryNode<T extends Comparable<T>> {

    //節點資料
    @Setter@Getter
    private T element;
    //左子節點
    @Setter@Getter
    private BinaryNode<T> left;
    //右子節點
    @Setter@Getter
    private BinaryNode<T> right;
    //節點高度
    @Setter@Getter
    private Integer height;

    //建構函式
    public BinaryNode(T element) {
        if (element == null) {
            throw new RuntimeException("二叉樹節點元素不能為空");
        }
        this.element = element;
        this.height = 0;
    }
}

我們現在改造BinaryNode類，並在類中增加高度屬性，高度預設為0。

/**
 * 二叉查詢樹
 */
public class BinarySearchTree<T extends Comparable<T>> {
    ……
    /**
     * 插入元素
     *
     * @param element
     */
    public void insert(T element) {
        root = insert(root, element);
    }

    private BinaryNode<T> insert(BinaryNode<T> tree, T element) {
        if (tree == null) {
            tree = new BinaryNode<>(element);
        } else {
            int compareResult = element.compareTo(tree.getElement());
            if (compareResult > 0) {
                tree.setRight(insert(tree.getRight(), element));
            }

            if (compareResult < 0) {
                tree.setLeft(insert(tree.getLeft(), element));
            }
        }

        return balance(tree);
    }
    
    /**
     * 平衡節點
     * @param tree
     */
    private BinaryNode<T> balance(BinaryNode<T> tree) {
        if (tree == null) {
            return null;
        }
        Integer leftHeight = height(tree.getLeft());
        Integer rightHeight = height(tree.getRight());
        if (leftHeight - rightHeight > 1) {
            //左-左情形，單旋轉
            if (height(tree.getLeft().getLeft()) >= height(tree.getLeft().getRight())) {
                tree = rotateWithLeftChild(tree);
            }
        } 

        //當前節點的高度 = 最高的子節點 + 1
        tree.setHeight(Math.max(leftHeight,rightHeight) + 1);
        return tree;
    }
        
	/**
     * 節點的高度
     * @param node
     * @return
     */
    public Integer height(BinaryNode node) {
        return node == null?-1:node.getHeight();
    }   
    
    /**
     * 左側單旋轉
     * @param k1
     */
    private BinaryNode<T> rotateWithLeftChild(BinaryNode<T> k1) {
        BinaryNode<T> k2 = k1.getLeft();
        k1.setLeft(k2.getRight());
        k2.setRight(k1);
        k1.setHeight(Math.max(height(k1.getLeft()),height(k1.getRight()))+1);
        k2.setHeight(Math.max(height(k2.getLeft()),height(k2.getRight()))+1);
        return k2;
    }
    ……
}

我們再在BinarySearchTree類中增加height方法，獲取節點的高度，如果節點為空，返回-1。由於insert後，樹可能會發生旋轉，節點會發生變化，所以這裡，insert方法改造為會有返回值。在第一個insert方法中，呼叫第二個insert方法，並用root去接第二個insert方法的返回值，說明整棵樹的根節點可能會發生旋轉變化。同樣在第二個insert方法中，遞迴呼叫時，根據不同的條件，將返回值給到當前節點的左或右子節點。節點插入完成後，我們統一呼叫balance方法，如果節點不滿足平衡條件，我們要進行相應的旋轉，最後把相關的節點的高度進行更新，這個balance方法是我們今天重點的方法。

進入balance方法後，我們分別獲取左右子樹的高度，如果左子樹的高度比右子樹高度大於1，說明不滿足平衡條件，需要進行旋轉。然後再判斷左子樹的左子樹與左子樹的右子樹的高度，如果大於，說明是左-左情形，需要左側單旋轉。這裡比較繞，大家多看幾篇，加深理解。我們把以當前節點為根節點的子樹傳入rotateWithLeftChild方法中，為了和上面的圖對應起來，變數的名稱叫做k1。那麼對應的k2就是k1的左子樹，然後進行旋轉，k1的左子樹設定為k2的右子樹，k2的右子樹設定為k1，然後再重新計算k1和k2的高度，最後將k2作為新子樹的根節點返回。這樣左-左情形的單旋轉就實現了。我們可以多看幾遍程式碼加深一下理解。

右右單旋轉

與左左相對稱的是右-右情形，我們看下圖：

我們插入節點6後，導致以k1為根節點的子樹不平衡，需要進行旋轉，旋轉的動作與左左情形完全對稱，總結操作如下：

  1. 發現k1的右子樹比左子樹的高度大於1；
  2. 發現k1的右子樹k2的右子樹高度大於k2的左子樹高度，這種稱作右-右情形。要做右側單旋轉。
  3. 將k2作為新樹的節點，k2的左子樹改為k1，k1的右子樹改為k2的左子樹。
  4. 更新k1和k2的高度。

旋轉後，如下圖：

我們按照上面的操作進行編碼，

/**
 * 平衡節點
 * @param tree
 */
private BinaryNode<T> balance(BinaryNode<T> tree) {
    if (tree == null) {
        return null;
    }
    Integer leftHeight = height(tree.getLeft());
    Integer rightHeight = height(tree.getRight());
    if (leftHeight - rightHeight > 1) {
        //左-左情形，單旋轉
        if (height(tree.getLeft().getLeft()) >= height(tree.getLeft().getRight())) {
            tree = rotateWithLeftChild(tree);
        }
    } else if (rightHeight - leftHeight > 1){
        //右-右情形，單旋轉
        if (height(tree.getRight().getRight()) >= height(tree.getRight().getLeft())) {
            tree = rotateWithRightChild(tree);
        }
    }

    //當前節點的高度 = 最高的子節點 + 1
    tree.setHeight(Math.max(leftHeight,rightHeight) + 1);
    return tree;
}

/**
 * 右側單旋轉
 * @param k1
 * @return
 */
private BinaryNode<T> rotateWithRightChild(BinaryNode<T> k1) {
    BinaryNode<T> k2 = k1.getRight();
    k1.setRight(k2.getLeft());
    k2.setLeft(k1);
    k1.setHeight(Math.max(height(k1.getLeft()),height(k1.getRight()))+1);
    k2.setHeight(Math.max(height(k2.getLeft()),height(k2.getRight()))+1);

    return k2;
}

在balance方法中，我們增加了右-右情形的判斷，然後呼叫rotateWithRightChild方法，在這個方法中，為了和上圖對應，變數的名字我們依然叫做k1和k2。k1的右節點設定為k2的左節點，k2的左節點設定為k1，然後更新高度，最後把新的根節點k2返回。

左右雙旋轉

下面我們再看雙旋轉的情形，如下圖所示：

我們新插入節點3後，導致以k1為根節點的子樹不滿足平衡條件，我們先用之前的左側單旋轉，看看能不能滿足，如下圖所示：

旋轉後，以k2為根節點的新樹，右子樹比左子樹的高度大於1，也不滿足平衡條件，所以這種方案是不行的。那我們要怎麼做呢？我們只有將k3作為新的根節點才能滿足平衡條件，將k3移動到根節點我們需要旋轉兩次，第一次先在k2節點進行右旋轉，將k3旋轉到k1的左子節點的位置，如圖：

然後再在k1位置進行左旋轉，將k3移動到根節點，如圖：

這樣就滿足了平衡條件，細心的小夥伴可能注意到了，原k3的做節點掛到了k2的右節點上，原k3的右節點刮到了k1的左節點上。這些細節並不需要我們特殊的處理，因為在左旋轉右旋轉的方法中已經處理過了，我們再總結一下具體的細節：

  1. 插入節點後，發現k1的左子樹比右子樹高度大於1；
  2. 發現k1的左子樹k2，k2的右子樹比k2的左子樹高，這是左-右情形，需要雙旋轉。
  3. 將k1的左子樹k2進行右旋轉；
  4. 將k1進行左旋轉；

我們編碼實現

/**
 * 平衡節點
 * @param tree
 */
private BinaryNode<T> balance(BinaryNode<T> tree) {
    if (tree == null) {
        return null;
    }
    Integer leftHeight = height(tree.getLeft());
    Integer rightHeight = height(tree.getRight());
    if (leftHeight - rightHeight > 1) {
        //左-左情形，單旋轉
        if (height(tree.getLeft().getLeft()) >= height(tree.getLeft().getRight())) {
            tree = rotateWithLeftChild(tree);
        } else {// 左-右情形，雙旋轉
            tree = doubleWithLeftChild(tree);
        }
    } else if (rightHeight - leftHeight > 1){
        //右-右情形，單旋轉
        if (height(tree.getRight().getRight()) >= height(tree.getRight().getLeft())) {
            tree = rotateWithRightChild(tree);
        }
    }

    //當前節點的高度 = 最高的子節點 + 1
    tree.setHeight(Math.max(leftHeight,rightHeight) + 1);
    return tree;
}

/**
 * 左側雙旋轉
 * @param k1
 * @return
 */
private BinaryNode<T> doubleWithLeftChild(BinaryNode<T> k1) {
    k1.setLeft(rotateWithRightChild(k1.getLeft()));
    return rotateWithLeftChild(k1);
}

我們在balance方法中，增加左-右情形的判斷，然後呼叫doubleWithLeftChild方法，在這個方法中，我們按照之前總結的步驟，先將k1的左節點進行一次右旋轉，然後再將k1進行左旋轉，最後將新的根節點返回，旋轉後達到了平衡的條件。

右左雙旋轉

最後我們再來看與左右情形對稱的右-左情形，樹的初始結構如下圖：

插入節點8後，導致k1節點的右子樹高度比左子樹高度大於1，同時k2的左子樹比右子樹高，這就是右-左情形。這時，我們需要先在k2節點做一次左旋轉，旋轉後如圖：

然後再在k1節點做一次右旋轉，旋轉後如圖：

我們參照上面的左右情形，總結一下右左情形的操作：

  1. 插入節點後，發現k1的右子樹比左子樹高度大於1；
  2. 發現k1的右子樹k2，k2的左子樹比k2的右子樹高，這是右-左情形，需要雙旋轉。
  3. 將k1的右子樹k2進行左旋轉；
  4. 將k1進行右旋轉；

然後我們編碼實現：

/**
 * 平衡節點
 * @param tree
 */
private BinaryNode<T> balance(BinaryNode<T> tree) {
    if (tree == null) {
        return null;
    }
    Integer leftHeight = height(tree.getLeft());
    Integer rightHeight = height(tree.getRight());
    if (leftHeight - rightHeight > 1) {
        //左-左情形，單旋轉
        if (height(tree.getLeft().getLeft()) >= height(tree.getLeft().getRight())) {
            tree = rotateWithLeftChild(tree);
        } else {// 左-右情形，雙旋轉
            tree = doubleWithLeftChild(tree);
        }
    } else if (rightHeight - leftHeight > 1){
        //右-右情形，單旋轉
        if (height(tree.getRight().getRight()) >= height(tree.getRight().getLeft())) {
            tree = rotateWithRightChild(tree);
        } else {//右-左情形，雙旋轉
            tree = doubleWithRightChild(tree);
        }
    }

    //當前節點的高度 = 最高的子節點 + 1
    tree.setHeight(Math.max(leftHeight,rightHeight) + 1);
    return tree;
}

/**
 * 右側雙旋轉
 * @param k1
 * @return
 */
private BinaryNode<T> doubleWithRightChild(BinaryNode<T> k1) {
    k1.setRight(rotateWithLeftChild(k1.getRight()));
    return rotateWithLeftChild(k1);
}

由於左右單旋轉的方法在之前已經實現過了，所以雙旋轉的實現，我們直接呼叫就可以了，先將k1的右節點進行一次左旋轉，再將k1進行右旋轉，最後返回新的根節點。因為節點的高度正在左右單旋轉的方法裡已經處理了，所以這裡不需要特殊的處理。

刪除節點

與插入節點一樣，刪除節點也會引起樹的不平衡，同樣，在刪除節點後，我們呼叫balance方法使樹再平衡。remove改造方法如下：

/**
 * 刪除元素
 * @param element
 */
public void remove(T element) {
    root = remove(root, element);
}

private BinaryNode<T> remove(BinaryNode<T> tree, T element) {
    if (tree == null) {
        return null;
    }
    int compareResult = element.compareTo(tree.getElement());
    if (compareResult > 0) {
        tree.setRight(remove(tree.getRight(), element));
    } else if (compareResult < 0) {
        tree.setLeft(remove(tree.getLeft(), element));
    }
    if (tree.getLeft() != null && tree.getRight() != null) {
        tree.setElement(findMin(tree.getRight()));
        tree.setRight(remove(tree.getRight(), tree.getElement()));
    } else {
        tree = tree.getLeft() != null ? tree.getLeft() : tree.getRight();
    }
    return balance(tree);
}

同樣，remove方法會引起子樹根節點的變化，所以，第二個remove方法要增加返回值，在呼叫第二個remove方法時，要用返回值覆蓋當前的節點。

總結

好了，AVL平衡二叉樹的操作就完全實現了，它解決了樹的不平衡問題，使得查詢效率大幅提升。小夥伴們有問題，歡迎評論區留言~~

相關文章