[算法]AVL树

来源：互联网发布：mac桌面壁纸不能平铺编辑：程序博客网时间：2024/06/04 19:27
#include <iostream>using namespace std;struct TreeNode {    int key;    int height{0};    TreeNode* left, * right;    TreeNode(int tKey = 0, TreeNode* tLeft = nullptr,             TreeNode* tRight = nullptr) {        key = tKey;        left = tLeft;        right = tRight;    }};class AVLTree {private:    TreeNode* root;    TreeNode* insertNodeImpl(int key, TreeNode* root) {        if (!root) {            root = new TreeNode(key);        } else {            if (key < root->key) {                root->left = insertNodeImpl(key, root->left);                root = fixUpWhenInsertingIntoLeft(key, root);            } else {                root->right = insertNodeImpl(key, root->right);                root = fixUpWhenInsertingIntoRight(key, root);            }        }        updateHeight(root);        return root;    }    int height(TreeNode* root) {        if (!root) {            return 0;        } else {            return root->height;        }    }    void updateHeight(TreeNode* root) {        root->height = max(height(root->left), height(root->right)) + 1;    }    TreeNode* fixUpWhenInsertingIntoLeft(int key, TreeNode* root) {        if (height(root->left) - height(root->right) >= 2) {            if (key < root->left->key) {                root = singleRotateWithLeft(root);            } else {                root = doubleRotateWithLeft(root);            }        }        return root;    }    TreeNode* fixUpWhenInsertingIntoRight(int key, TreeNode* root) {        if (height(root->right) - height(root->left) >= 2) {            if (key > root->right->key) {                root = singleRotateWithRight(root);            } else {                root = doubleRotateWithRight(root);            }        }        return root;    }    TreeNode* singleRotateWithLeft(TreeNode* root) {        TreeNode* newRoot = root->left;        root->left = newRoot->right;        newRoot->right = root;        updateHeight(root);        updateHeight(newRoot);        return newRoot;    }    TreeNode* singleRotateWithRight(TreeNode* root) {        TreeNode* newRoot = root->right;        root->right = newRoot->left;        newRoot->left = root;        updateHeight(root);        updateHeight(newRoot);        return newRoot;    }    TreeNode* doubleRotateWithLeft(TreeNode* root) {        root->left = singleRotateWithRight(root->left);        return singleRotateWithLeft(root);    }    TreeNode* doubleRotateWithRight(TreeNode* root) {        root->right = singleRotateWithLeft(root->right);        return singleRotateWithRight(root);    }    void traversalImpl(TreeNode* root) {        if (root) {            cout << root->key << " ";            traversalImpl(root->left);            traversalImpl(root->right);        }    }    TreeNode* deleteNodeImpl(int key, TreeNode* root) {        if (!root) { return root; }        if (key == root->key) {            if (!root->right) {                TreeNode* newRoot = root->left;                delete root;                return newRoot;            } else {                int min = findMin(root->right);                root->key = min;                root->right = deleteNodeImpl(min, root->right);            }        } else if (key > root->key) {            root->right = deleteNodeImpl(key, root->right);        } else {            root->left = deleteNodeImpl(key, root->left);        }        updateHeight(root);        root = deleteFixUp(root);        return root;    }    int findMin(TreeNode* root) {        while (root->left) {            root = root->left;        }        return root->key;    }    TreeNode* deleteFixUp(TreeNode* root) {        if (height(root->left) - height(root->right) >= 2) {            if (height(root->left->left) > height(root->left->right)) {                root = singleRotateWithLeft(root);            } else {                root = doubleRotateWithLeft(root);            }        } else if (height(root->right) - height(root->left) >= 2) {            if (height(root->right->right) > height(root->right->left)) {                root = singleRotateWithRight(root);            } else {                root = doubleRotateWithRight(root);            }        }        return root;    }public:    AVLTree() { root = nullptr; }    void insertNode(int key) {        root = insertNodeImpl(key, root);    }    void deleteNode(int key) {        root = deleteNodeImpl(key, root);        updateHeight(root);        root = deleteFixUp(root);    }    void traversal() {        traversalImpl(root);        cout << endl;    }};int main() {    AVLTree tree;    for (int i = 0; i < 4; i += 1) {        tree.insertNode(i);    }    tree.insertNode(19);    tree.traversal();    tree.deleteNode(15);    tree.traversal();    return 0;}
0 0