手寫AVL平衡二叉搜尋樹

二叉搜尋樹的侷限性

先說一下什麼是二叉搜尋樹，二叉樹每個節點只有兩個節點，二叉搜尋樹的每個左子節點的值小於其父節點的值，每個右子節點的值大於其左子節點的值。如下圖：

二叉搜尋樹，顧名思義，它的搜尋效率很高，可以達到O(logn)。但這是理想狀況下的，即上圖所示。實際上，由於插入順序的原因，形成的二叉搜尋樹並不會像上圖這樣“工整”，最壞的情況的下，甚至可能會退化成連結串列了，如下圖：

這顯然不是我們想要看的結果，那麼我們必須要引入一套機制來避免這種事情的發生，也就是讓二叉搜尋樹帶上平衡條件。

AVL平衡二叉搜尋樹

幾個基本概念

  • 葉子節點：既沒有左子節點，也沒有右左子節點的節點就是葉子節點。

  • 樹的高度：葉子節點的高度為1，空節點的高度是-1，父節點的高度是其兩個子樹較高一棵子樹的高度加一。

  • 平衡條件：每一個節點的左子樹與右子樹的高度差不超過1。

核心思想

因為AVL平衡二叉搜尋樹，父節點的兩顆子樹的高度差不能超過1。在AVL平衡二叉樹種，採用旋轉的機制來使不滿足平衡條件的二叉樹重新回到滿足平衡條件的狀態。在二叉搜尋樹中，需要被平衡的情況可以分為兩大類總共四種情況，

  • 單旋轉

    • 左旋轉
    • 右旋轉
  • 雙旋轉
    • 先左旋，再右旋
    • 先右旋，再左旋

如下圖所示：

通過圖片的形式我們很容易就可以寫出使二叉樹回到滿足平衡條件的程式碼

  1. // 右旋轉
  2. public TreeNode rightRotate(TreeNode root) {
  3.         TreeNode temp1 = root.left;
  4.         TreeNode temp2 = temp1.right;
  5.         temp1.right = root;
  6.         root.left = temp2;
  7.         return temp1;
  8. }
  10. // 左旋轉
  11. public TreeNode leftRotate(TreeNode root) {
  12.     TreeNode temp1 = root.right;
  13.     TreeNode temp2 = temp1.left;
  14.     temp1.left = root;
  15.     root.right = temp2;
  16.     return temp1;
  17. } 
  19. // 先右後左
  20. public TreeNode rightLeftRotate(TreeNode root) {
  21.     root.right = rightRotate(root.right);
  22.     return leftRotate(root);
  23. }
  25. // 先左後右
  26. public TreeNode leftRightRotate(TreeNode root) {
  27.     root.left = leftRotate(root.left);
  28.     return rightRotate(root);
  29. }

我們必須再每一次插入節點後判斷樹是否需要平衡，也就是是否會出現兩顆子樹的高度差超過1的情況，首先編寫一個可以計算出傳入節點 高度的函式。

  1. public int height(TreeNode root) {
  2.     if (root == null) {
  3.         return -1;
  4.     }
  5.     if (root.left == null && root.right == null) {
  6.         return 0;
  7.     }
  8.     return Math.max(height(root.right), height(root.left));
  9. }

有了這個函式，我們就不僅可以判斷是否出現需要平衡的情況，還可以判斷需要平衡的情況是四種情況種的哪一種。

  1. public TreeNode balance(TreeNode root) {
  2.     int l = height(root.left);
  3.     int r = height(root.right);
  4.     if (l - r >= 2) {
  5.         // rightRotate
  6.         if (height(root.left.left) - height(root.left.right) >= 1) {
  7.             // rightRotate
  8.             root = rightRotate(root);
  9.         } else if (height(root.left.right) - height(root.left.left) >= 1) {
  10.             // leftRightRotate
  11.             root = leftRightRotate(root);
  12.         }
  13.     } else if (r - l >= 2) {
  14.         // leftRotate
  15.         if (height(root.right.right) - height(root.right.left) >= 1) {
  16.             // leftRotate
  17.             root = leftRotate(root);
  18.         } else if (height(root.right.left) - height(root.right.right) >= 1){
  19.             root = rightLeftRotate(root);
  20.         }
  21.     }
  22.     return root;
  23. }

以上就是AVL平衡二叉搜尋樹的精髓，並且已經用程式碼實現了。

完整程式碼

這是我完善後的功能相對完整的AVL二叉搜尋平衡樹。

  1. class TreeNode {
  2.     int value;
  3.     TreeNode left;
  4.     TreeNode right;
  5.     int height;
  7.     public TreeNode(int value) {
  8.         this.value = value;
  9.     }
  10. }
  12. public class AVLBinarySearchTree {
  14.     public static void main(String[] args) {
  15.         AVLBinarySearchTree a = new AVLBinarySearchTree();
  16.         for (int i = 0; i < 10; i++) {
  17.             a.insert(i);
  18.         }
  19.     }
  21.     private TreeNode root;
  22.     private static final int ALLOWED_IMBALANCE = 1;
  24.     // 刪除元素
  25.     public void remove(int value) {
  26.         root = remove(value, root);
  27.     }
  29.     // 檢查是否包含某一元素，包含則返回該節點，不包含則返回null
  30.     public TreeNode contain(int value) {
  31.         TreeNode temp = root;
  32.         if (temp == null) {
  33.             return temp;
  34.         }
  35.         while (temp.value != value) {
  36.             if (value > temp.value) {
  37.                 temp = temp.right;
  38.             } else if (value < temp.value) {
  39.                 temp = temp.left;
  40.             }
  41.         }
  42.         return temp;
  43.     }
  45.     // 刪除指定子樹上的指定元素
  46.     private TreeNode remove(int value, TreeNode abn) {
  47.         if (abn == null) {
  48.             return abn;
  49.         }
  51.         if (value > abn.value) {
  52.             abn.right = remove(value, abn.right);
  53.         } else if (value < abn.value) {
  54.             abn.left = remove(value, abn.left);
  55.         } else {
  56.             if (abn.right == null && abn.left == null) {
  57.                 abn = null;
  58.                 return abn;
  59.             } else if (abn.right != null) {
  60.                 abn.value = findMin(abn.right).value;
  61.                 abn.right = remove(abn.value, abn.right);
  62.             } else {
  63.                 abn.value = findMax(abn.left).value;
  64.                 abn.left = remove(abn.value, abn.left);
  65.             }
  67.         }
  69.         return balance(abn);
  70.     }
  72.     // 找到指定子樹最大值
  73.     private TreeNode findMax(TreeNode abn) {
  74.         if (abn == null) {
  75.             return null;
  76.         }
  77.         TreeNode temp = abn;
  78.         while (temp.right != null) {
  79.             temp = temp.right;
  80.         }
  81.         return temp;
  83.     }
  85.     // 找到指定子樹最小值
  86.     private TreeNode findMin(TreeNode abn) {
  87.         if (abn == null) {
  88.             return null;
  89.         }
  90.         TreeNode temp = abn;
  91.         while (temp.left != null) {
  92.             temp = temp.left;
  93.         }
  94.         return temp;
  95.     }
  97.     // 插入節點
  98.     public void insert(int value) {
  99.         root = insert(value, root);
  100.     }
  102.     // 計算節點高度
  103.     private int height(TreeNode abn) {
  104.         if (abn == null) {
  105.             return -1;
  106.         }
  107.         return abn.height;
  108.     }
  110.     // 樹的高度
  111.     public int height() {
  112.         return height(root);
  113.     }
  115.     // 插入節點
  116.     private TreeNode insert(int value, TreeNode abn) {
  117.         if (abn == null) {
  118.             return new TreeNode(value);
  119.         }
  120.         if (value > abn.value) {
  121.             abn.right = insert(value, abn.right);
  122.         } else if (value < abn.value) {
  123.             abn.left = insert(value, abn.left);
  124.         }
  125.         return balance(abn);
  126.     }
  128.     // 平衡不平衡的樹
  129.     private TreeNode balance(TreeNode abn) {
  130.         if (height(abn.left) - height(abn.right) > ALLOWED_IMBALANCE) {
  131.             if (height(abn.left.left) >= height(abn.left.right)) {
  132.                 abn = leftSingleRotate(abn);
  133.             } else if (height(abn.left.left) < height(abn.left.right)) {
  134.                 abn = leftDoubleRotate(abn);
  135.             }
  136.         } else if (height(abn.right) - height(abn.left) > ALLOWED_IMBALANCE) {
  137.             if (height(abn.right.right) >= height(abn.right.left)) {
  138.                 abn = rightSingleRotate(abn);
  139.             } else {
  140.                 abn = rightDoubleRotate(abn);
  141.             }
  142.         }
  143.         abn.height = Math.max(height(abn.left), height(abn.right)) + 1;
  144.         return abn;
  145.     }
  147.     // 右單旋轉
  148.     private TreeNode rightSingleRotate(TreeNode abn) {
  149.         TreeNode temp = abn;
  150.         abn = abn.right;
  151.         temp.right = abn.left;
  152.         abn.left = temp;
  153.         temp.height = Math.max(height(temp.right), height(temp.left)) + 1;
  154.         abn.height = Math.max(height(abn.right), temp.height) + 1;
  155.         return abn;
  156.     }
  158.     // 左單旋轉
  159.     private TreeNode leftSingleRotate(TreeNode abn) {
  160.         TreeNode temp = abn;
  161.         abn = abn.left;
  162.         temp.left = abn.right;
  163.         abn.right = temp;
  164.         temp.height = Math.max(height(temp.right), height(temp.left)) + 1;
  165.         abn.height = Math.max(height(abn.right), temp.height) + 1;
  166.         return abn;
  167.     }
  169.     // 右雙旋轉
  170.     private TreeNode rightDoubleRotate(TreeNode abn) {
  171.         abn.right = leftSingleRotate(abn.right);
  172.         return rightSingleRotate(abn);
  173.     }
  175.     // 左雙旋轉
  176.     private TreeNode leftDoubleRotate(TreeNode abn) {
  177.         abn.left = rightSingleRotate(abn.left);
  178.         return leftSingleRotate(abn);
  179.     }
  180. }