[C/C++] AVL樹，插入及刪除

#include <stdlib.h>
#include <iostream> 
#include <algorithm>

using namespace std;

class Node{
    private:
        int data;
        int height;
        Node *left, *right;
    public:
        Node(int val){ data = val; left = right = NULL; height = 1;}
        friend class AVLTree;
};

class AVLTree{
    private:
        int height(Node* node){
            return (!node)? 0: node->height;
        }
        int getBF(Node* node){
            return (!node)? 0: height(node->left) - height(node->right);
        }
        //旋轉方法
        Node* rotate(Node* old_rootT, char direct) {
            //重新定義此子樹的rootT節點
            Node *new_rootT = (direct=='L')? old_rootT->right: old_rootT->left;  //若左轉，新的rootT會是舊rootT->R，若右轉，新的rootT會是舊rootT->L
            Node *temp = (direct=='L')? new_rootT->left: new_rootT->right;       //若左轉，暫存子樹會是舊rootT->R->L（即新的rootT>L），若右轉，暫存子樹會是舊rootT->L->R（即新的rootT>R）

            //旋轉
            if(direct=='L'){                           //若左轉
                new_rootT->left = old_rootT;               //舊rooT會變成新的rootT->L
                old_rootT->right = temp;                   //舊rootT->R會是暫存子樹temp（即舊rootT->R->L, 亦即新的rootT->L)
            }else{                                     //若右轉
                new_rootT->right = old_rootT;              //舊rooT會變成新的rootT->R
                old_rootT->left = temp;                    //舊rootT->L會是暫存子樹temp（即舊rootT->L->R, 亦即新的rootT->R)
            }
        
            //更新有改變到的高度（從底至上），所以先update舊rootT
            old_rootT->height = max(height(old_rootT->left), height(old_rootT->right)) + 1;  
            new_rootT->height = max(height(new_rootT->left), height(new_rootT->right)) + 1;

            return new_rootT;
        }

        //插入遞迴式
        Node* insertRecursive(Node* rootT, int val){
            //使用一般的BST樹插入方法插入數字
            if(!rootT) return (new Node(val));                                              //若目前節點 (或作 目前子樹的根) 為空，即樹葉節點插入新節點至樹葉 (回傳給上一層，此例中上一層也等同rootT，因為是用 rootT = insert(rootT, val)呼叫 )
            else if(val < rootT->data) rootT->left = insertRecursive(rootT->left, val);     //若 數字 < 目前節點，則插入左子樹 (即：rootT->左子 = insert(rootT->左子) 會 return 一個插入完數字 的子樹 的rootT 給 rootT->左子)
            else if(val > rootT->data) rootT->right = insertRecursive(rootT->right, val);   //若 數字 > 目前節點，則插入右子樹
            else return rootT;                                                              //二元樹原則上非key-value數對型態者，不會有重複值
                                                                                            //把rootT 回傳給rootT (即：把自己丟回給自己，不作修改)
            
            rootT->height = max(height(rootT->left), height(rootT->right)) + 1;             //插入後，更新目前這個節點的樹高 (取左右子樹高之大者 + 1(自己))
            int bf = getBF(rootT);                                                          //取得平衡因子

            //process if abs(bf)>1
            if(bf > 1 && val < rootT->left->data){       //(平衡因子>1) => 左子樹高，前面必插左邊;   (插入數 < 左子樹值) => 前面必插左子樹的左邊 {LL}
                return rotate(rootT, 'R');                      //自己 (目前子樹) 右轉，然後把rootT 回傳給上一層 (同前面提到，此例中上一層也等同rootT，即原本呼叫這個函數的人rootT = insert(rootT, val) )
            }
            if(bf > 1 && val > rootT->left->data){       //(平衡因子>1) => 左子樹高，前面必插左邊;   (插入數 > 左子樹值) => 前面必插左子樹的右邊 {LR}
                rootT->left = rotate(rootT->left, 'L');         //自己的左邊 (目前子樹的左子樹) 先左轉
                return rotate(rootT, 'R');                      //自己 (目前子樹) 再右轉
            }
            if(bf < -1 && val > rootT->right->data){     //(平衡因子< -1) => 右子樹高，前面必插右邊;   (插入數 > 右子樹值) => 前面必插右子樹的右邊 {RR}
                return rotate(rootT, 'L');                      //自己 (目前子樹) 左轉
            }
            if(bf < -1 && val < rootT->right->data){     //(平衡因子< -1) => 右子樹高，前面必插右邊;   (插入數 < 右子樹值) => 前面必插右子樹的左邊 {RL}
                rootT->right = rotate(rootT->right, 'R');       //自己的右邊 (目前子樹的左子樹) 先右轉
                return rotate(rootT, 'L');                      //自己 (目前子樹) 再左轉
            }

            return rootT;
        }
        //找子樹中最小的樹 (從根開始往左走直到不能走)
        Node* minValue(Node* rootT){
            Node* ptr = rootT;
            while(ptr->left)
                ptr = ptr->left;
            return ptr;
        }
        //自AVL樹中刪除
        Node* removeRecursive(Node* rootT, int val){
            //使用一般的BST樹插入方法刪除數字
            if(!rootT) return NULL;                                                           //若為空值，那不用刪除，直接回傳給剛剛叫他的人(亦作 return rootT;) (rootT==NULL 即剛剛叫他的人傳入的子樹為空)
            else if(val < rootT->data) rootT->left = removeRecursive(rootT->left, val);       //若 數字 < 目前節點，要刪除的數字在目前子樹的左子樹
            else if(val > rootT->data) rootT->right = removeRecursive(rootT->right, val);     //若 數字 > 目前節點，要刪除的數字在目前子樹的右子樹
            else{                                                                             //上面都不成立，即找到相等的數，進入刪除階段
                if(!rootT->left || !rootT->right){                                                  //若目前節點的子樹數量為 0(無子樹) 或 只有單邊子樹
                    Node* childT = (rootT->left)? rootT->left: rootT->right;                                //暫存左或右子樹
                    delete rootT;                                                                           //釋放目前子樹樹根記憶體
                    rootT = childT;                                                                         //令左或右子樹樹根為樹根 (左或右樹直接上拉)
                }else{                                                                              //若目前節點 左、右子樹皆有
                    Node* minT = minValue(rootT->right);                                                    //取得 右子樹最小的節點
                    rootT->data = minT->data;                                                               //設目前子樹樹根的數字 為剛剛找到的節點的數字
                    rootT->right = removeRecursive(rootT->right, minT->data);                               //刪除剛剛找到的節點 (註：因為右子樹最小的節點必無左子樹，故遞迴結果必為第1或第2個if)
                }
            }
    
            if(!rootT) return rootT;                                                          //若目前子樹樹根為空，必平衡且樹高為0，直接回傳

            rootT->height = max(height(rootT->left), height(rootT->right)) + 1;               //更新目前這個節點的樹高
            int bf = getBF(rootT);                                                            //取得平衡因子

            //process if abs(bf)>1
            if(bf > 1 && getBF(rootT->left) >= 0){          //(平衡因子>1) => 左子樹高，前面必刪右邊;   (左子節點平衡因子>=0的話) => 右轉後必變為-1或0，故不用先轉; {LL}
                return rotate(rootT, 'R');                          //自己 (目前子樹) 右轉
            }
            if(bf > 1 && getBF(rootT->left) < 0){           //(平衡因子>1) => 左子樹高，前面必刪右邊;   (左子節點平衡因子<0的話) => 右轉後必變為-2，故要先轉; {LR}
                rootT->left = rotate(rootT->left, 'L');             //自己的左邊 (目前子樹的左子樹) 先左轉
                return rotate(rootT, 'R');                          //自己 (目前子樹) 再右轉
            }
            if(bf < -1 && getBF(rootT->right) <= 0){        //(平衡因子< -1) => 右子樹高，前面必刪左邊;  (右子節點平衡因子<=0的話) => 左轉後必變為1或0，故不用先轉; {RR}
                return rotate(rootT, 'L');                          //自己 (目前子樹) 左轉
            }
            if(bf < -1 && getBF(rootT->right) > 0){         //(平衡因子< -1) => 右子樹高，前面必刪左邊;  (右子節點平衡因子>0的話) => 左轉後必變為2，故要先轉; {RL}
                rootT->right = rotate(rootT->right, 'R');           //自己的右邊 (目前子樹的右子樹) 先右轉
                return rotate(rootT, 'L');                          //自己 (目前子樹) 再左轉
            }

            return rootT;                            
        }
    public:
        Node *root;

        AVLTree(){ root = NULL; }

        void insert(int val) {
            root = insertRecursive(root, val);
        }

        void remove(int val) {
            root = removeRecursive(root, val);
        }
        //前序印出AVL樹
        void preOrder(Node* root) {
            if(root) {
                cout << root->data << "\t";
                root->left? (cout << root->left->data << "\t"): (cout << "Nan\t");
                root->right? (cout << root->right->data << endl): (cout << "Nan" << endl);
                preOrder(root->left);
                preOrder(root->right);
            }
        }
};

int main(void){
    AVLTree AVLTree;
    int val;
    char SENT;

    //對AVL樹的操作
    cout << "<1: 插入> <2: 刪除> <3: 列出(前序)> <q: 離開>" << endl << "選擇操作 : ";
    cin >> SENT;
    while(SENT != 'q'){
        switch(SENT) { 
            case '1':
                cout << "輸入要插入的數: ";
                cin >> val;
                AVLTree.insert(val);
                break; 
            case '2':
                cout << "輸入要刪除的數: ";
                cin >> val;
                AVLTree.remove(val);
                break; 
            case '3':
                cout << endl << "樹根節點按前序追蹤排序" << endl;
                cout << "樹根\t左子\t右子" << endl;
                AVLTree.preOrder(AVLTree.root);
                cout << endl;
                break; 
            default: 
                break;
        }
        cout << endl << "<1: 插入> <2: 刪除> <3: 列出(前序)> <q: 離開>" << endl << "選擇操作 : ";
        cin >> SENT;
    }

    system("pause");
    return 0;
}