AVL樹( C++)

#include <iostream>
#include <queue>
using namespace std;

struct Node {
    int val;
    Node *left = nullptr;
    Node *right = nullptr;
    int factor = 0;

    Node(int i) : val(i) {};
};

class AVLTree {
public:
    void balance(Node *&node) {
        if (node->factor == 2) {
            if (node->left->factor == 1) {
                Node *temp = node->left;
                node->left = temp->right;
                node->factor = 0;
                temp->right = node;
                temp->factor = 0;
                node = temp;
            } else if (node->left->factor == -1) {
                Node *temp = node->left;
                node->left = temp->right;
                temp->right = nullptr;
                node->left->left = temp;
                node->left->factor = 1;
                if (node->left->right == nullptr) {
                    node->left->left->factor = 0;
                } else {
                    node->left->left->factor = 1;
                }
                balance(node);
            }
        } else if (node->factor == -2) {
            if (node->right->factor == -1) {
                Node *temp = node->right;
                node->right = temp->left;
                node->factor = 0;
                temp->left = node;
                temp->factor = 0;
                node = temp;
            } else if (node->right->factor == 1) {
                Node *temp = node->right;
                node->right = temp->left;
                temp->left = nullptr;
                node->right->right = temp;
                node->right->factor = -1;
                if (node->right->left == nullptr) {
                    node->right->right->factor = 0;
                } else {
                    node->right->right->factor = -1;
                }
                balance(node);
            }
        }
    }

    void insert(Node *&node, int val, bool &adjust) {
        if (node == nullptr) return;
        if (val < node->val) {
            if (node->left == nullptr) {
                Node *temp = new Node(val);
                node->left = temp;
                ++node->factor;
                //高度不變,左右平衡
                if (node->factor == 0) adjust = false;
            } else {
                insert(node->left, val, adjust);
                //回溯,調整平衡因子
                if (adjust) ++node->factor;
            }
        } else if (val > node->val) {
            if (node->right == nullptr) {
                Node *temp = new Node(val);
                node->right = temp;
                --node->factor;
                //高度不變,左右平衡
                if (node->factor == 0) adjust = false;
            } else {
                insert(node->right, val, adjust);
                //回溯,調整平衡因子
                if (adjust) --node->factor;
            }
        } else {
            perror("same value");
            exit(1);
        }
        if (node->factor == 2 || node->factor == -2) {
            balance(node);
            //不管是LL還是LR全部旋轉之後的高度都是不變的,可畫圖驗證
            adjust = false;
        }
    }

    void insert(int val) {
        bool adjust = true;
        insert(root, val, adjust);
    }

    bool remove(Node *&node, int val, bool &shorter) {
        if (node == nullptr) return false;
        if (val < node->val) {
            if (remove(node->left, val, shorter)) {
                return false;
            }
            if (shorter) {
                switch (node->factor) {
                    case 1:
                        node->factor = 0;
                        shorter = true;
                        break;
                    case -1:
                        shorter = node->right->factor != 0;
                        node->factor = -2;
                        balance(node);
                        break;
                    case 0:
                        node->factor = -1;
                        shorter = false;
                        break;
                    default:
                        break;
                }
            }
        } else if (val > node->val) {
            if (remove(node->right, val, shorter)) {
                return false;
            }
            if (shorter) {
                switch (node->factor) {
                    case 1:
                        shorter = node->left->factor != 0;
                        node->factor = 2;
                        balance(node);
                        break;
                    case -1:
                        node->factor = 0;
                        shorter = true;
                        break;
                    case 0:
                        node->factor = 1;
                        shorter = false;
                        break;
                    default:
                        break;
                }
            }
        } else {
            if (node->left == nullptr) {
                Node *temp = node;
                node = node->right;
                delete temp;
                shorter = true;
            } else if (node->right == nullptr) {
                Node *temp = node;
                node = node->left;
                delete temp;
                shorter = true;
            } else {
                Node *temp = node->right;
                while (temp->left != nullptr) {
                    temp = temp->left;
                }
                node->val = temp->val;
                remove(node->right, temp->val, shorter);
                if (shorter) {
                    switch (node->factor) {
                        case 1:
                            shorter = node->left->factor != 0;
                            node->factor = 2;
                            balance(node);
                            break;
                        case -1:
                            node->factor = 0;
                            shorter = true;
                            break;
                        case 0:
                            node->factor = 1;
                            shorter = false;
                            break;
                        default:
                            break;
                    }
                }
            }
        }
        return true;
    }

    void remove(int val) {
        bool shorter = false;
        remove(root, val, shorter);
    }

    Node *create(int arr[], int len) {
        for (int i = 0; i < len; ++i) {
            if (this->root == nullptr) {
                root = new Node(arr[i]);
            } else {
                insert(arr[i]);
            }
        }
        return root;
    }

    void level_travel() {
        queue<Node *> q;
        q.push(root);
        while (!q.empty()) {
            Node *n = q.front();
            q.pop();
            cout << n->val << " ";
            if (n->left != nullptr) {
                q.push(n->left);
            }
            if (n->right != nullptr) {
                q.push(n->right);
            }
        }
        cout << "\n";
    }

private:
    Node *root = nullptr;
};

int main() {
    int len = 10;
    int nodes[10] = { 5,13,7,9,1,3,10,15,20,25 };
    AVLTree avl;
    avl.create(nodes, len);
    avl.level_travel();
    avl.remove(15);
    avl.level_travel();
    return 0;
}

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章