AVL树实现

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/hzh_0000/article/details/80044799

“AVLTree.h”

#pragma once
#include <algorithm>
#include <string>
#include <iostream>
using namespace std;


class AVLNode {
    //树节点
    friend class AVLTree;
    AVLNode *lchild, *rchild;
    int height;
    int key;

    AVLNode() {
        lchild = nullptr;
        rchild = nullptr;
        height = 1;
        key = 0;
    }

    AVLNode( int k) {
        lchild = nullptr;
        rchild = nullptr;
        height = 1;
        key = k;
    }
    ~AVLNode() {}
};

class AVLTree {
    AVLNode *root;

    /*
        向树中添加节点
    */
    AVLNode* _add_node(AVLNode* root, int key) {
        //如果是空节点，说明应该在该位置插入节点
        if (root == nullptr) {
            root = new AVLNode(key);
        }
        //如果比当前值小，则向左子树插入
        else if (key < root->key) {
            //向左子树插入，并更新左子树
            root->lchild = _add_node(root->lchild, key);

            //判断左子树是否打破平衡
            if (GetHeight(root->lchild) - GetHeight(root->rchild) == 2) {

                //如果是向左子树的左子树插入了节点，那么应该进行LL操作（对当前节点执行右旋）
                if (GetHeight(root->lchild->lchild) > GetHeight(root->lchild->rchild)) {
                    root = LL(root);
                }
                //如果是向左子树的右子树插入了节点，那么应该进行LR操作（对当前节点的左子树执行左旋，然后对当前节点执行右旋）
                else {
                    root = LR(root);
                }
            }

            //更新节点高度
            UpdateHeight(root);
        }
        //如果比当前值大，则向右子树插入
        else if (key > root->key) {
            //向右子树插入，并更新右子树
            root->rchild = _add_node(root->rchild, key);

            //判断右子树是否打破平衡
            if (GetHeight(root->lchild) - GetHeight(root->rchild) == -2) {

                //如果是向右子树的左子树插入了节点，那么应该进行RL操作（对当前节点的右子树执行右旋，然后对当前节点执行左旋）
                if (GetHeight(root->rchild->lchild) > GetHeight(root->rchild->rchild)) {
                    root = RL(root);
                }
                //如果是向右子树的右子树插入了节点，那么应该进行RR操作（对当前节点执行左旋）
                else {
                    root = RR(root);
                }
            }

            //更新节点高度
            UpdateHeight(root);
        }

        return root;
    }

    /*
    向树中删除节点
    */
    AVLNode* _delete_node(AVLNode* root, int key) {
        if (root == nullptr) {
            return nullptr;
        }
        if (root->key == key) {
            if (root->lchild == nullptr) {
                AVLNode* temp = root;
                root = root->rchild;
                delete temp;
            }
            else if (root->rchild == nullptr) {
                AVLNode* temp = root;
                root = root->lchild;
                delete temp;
            }
            else {
                //判断左子树高还是右子树高
                //如果左子树高，那么把当前节点的前驱提上来，然后在左子树中递归删除当前节点前驱
                if (GetHeight(root->lchild) > GetHeight(root->rchild)) {
                    AVLNode *temp = root->lchild, *pre = nullptr;
                    while (temp != nullptr) {
                        pre = temp;
                        temp = temp->rchild;
                    }
                    root->key = pre->key;
                    root->lchild = _delete_node(root->lchild, pre->key);
                }
                //如果右子树高，那么把当前节点的后驱提上来，然后在右子树中递归删除当前节点后驱
                else {
                    AVLNode *temp = root->rchild, *pre = nullptr;
                    while (temp != nullptr) {
                        pre = temp;
                        temp = temp->lchild;
                    }
                    root->key = pre->key;
                    root->rchild = _delete_node(root->rchild, pre->key);
                }
                UpdateHeight(root);
            }
        }
        else if (key < root->key) {
            root->lchild = _delete_node(root->lchild, key);

            //判断右子树是否打破平衡
            if (GetHeight(root->lchild) - GetHeight(root->rchild) == -2) {

                //如果是向右子树的左子树插入了节点，那么应该进行RL操作（对当前节点的右子树执行右旋，然后对当前节点执行左旋）
                if (GetHeight(root->rchild->lchild) > GetHeight(root->rchild->rchild)) {
                    root = RL(root);
                }
                //如果是向右子树的右子树插入了节点，那么应该进行RR操作（对当前节点执行左旋）
                else {
                    root = RR(root);
                }
            }

            //更新节点高度
            UpdateHeight(root);
        }
        else {
            root->rchild = _delete_node(root->rchild, key);

            //判断左子树是否打破平衡
            if (GetHeight(root->lchild) - GetHeight(root->rchild) == 2) {

                //如果是向左子树的左子树插入了节点，那么应该进行LL操作（对当前节点执行右旋）
                if (GetHeight(root->lchild->lchild) > GetHeight(root->lchild->rchild)) {
                    root = LL(root);
                }
                //如果是向左子树的右子树插入了节点，那么应该进行LR操作（对当前节点的左子树执行左旋，然后对当前节点执行右旋）
                else {
                    root = LR(root);
                }
            }

            //更新节点高度
            UpdateHeight(root);
        }
        return root;
    }


    //返回当前节点的高度
    int GetHeight(AVLNode *root) {
        return root == nullptr ? 0 : root->height;
    }

    //LL操作（对当前节点执行右旋）
    AVLNode* LL(AVLNode* root) {
        root = RRoate(root);
        return root;
    }

    //LR操作（对当前节点的左子树执行左旋，然后对当前节点执行右旋）
    AVLNode* LR(AVLNode* root) {
        root->lchild = LRoate(root->lchild);
        root = RRoate(root);
        return root;
    }

    //RL操作（对当前节点的右子树执行右旋，然后对当前节点执行左旋）
    AVLNode* RL(AVLNode* root) {
        root->rchild = RRoate(root->rchild);
        root = LRoate(root);
        return root;
    }

    //RR操作（对当前节点执行左旋）
    AVLNode* RR(AVLNode* root) {
        root = LRoate(root);
        return root;
    }

    //左旋
    AVLNode* LRoate(AVLNode* root) {
        AVLNode* p = root->rchild;
        root->rchild = p->lchild;
        p->lchild = root;

        UpdateHeight(root);
        UpdateHeight(p);

        return p;
    }

    //右旋
    AVLNode* RRoate(AVLNode* root) {
        AVLNode* p = root->lchild;
        root->lchild = p->rchild;
        p->rchild = root;

        UpdateHeight(root);
        UpdateHeight(p);

        return p;
    }

    //更新节点高度
    void UpdateHeight(AVLNode *root) {
        root->height = max(GetHeight(root->lchild), GetHeight(root->rchild)) + 1;
    }

    //打印这颗树
    void ppp(AVLNode* root, string s) {
        if (root == nullptr) {
            cout << endl;
            return;
        }
        ppp(root->lchild, s + "\t");
        cout << s << root->key << endl;
        ppp(root->rchild, s + "\t");
    }
public:
    void add_node(int key) {
        root = _add_node(root, key);
    }

    void delete_node(int key) {
        root = _delete_node(root, key);
    }

    AVLTree(): root(nullptr){}

    int GetTreeHeight() {
        return GetHeight(root);
    }

    //打印这颗树
    void pp() {
        ppp(root,"");
    }
};

测试：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<cstdlib>
#include<string>
#include<vector>
#include<stack>
#include<set>
#include<queue>
#include<map>
#include <random>
#include "AVLTree.h"
#include <memory>
#include <time.h>

using namespace std;



int main() {
    AVLTree tree;
    double sum = 0;
    for (int i = 0; i < 100000; i++) {
        int x = rand();
        time_t st = clock();
        tree.add_node(x);
        time_t ed = clock();
        sum += ed - st;
        if (i % 10000==0 && i != 0) {

            cout << x << "\t\t" << log(i) << "\t\t" << tree.GetTreeHeight() << "\t\t" << sum << endl;
            sum = 0;
        }

    }
    cout << "-----------------" << endl;
    for (int i = 0; i < 100000; i++) {
        int x = rand();
        time_t st = clock();
        tree.delete_node(x);
        time_t ed = clock();
        sum += ed - st;
        if (i % 10000 == 0 && i != 0) {

            cout << x << "\t\t" << log(i) << "\t\t" << tree.GetTreeHeight() << "\t\t" << sum << endl;
            sum = 0;
        }

    }
    return 0;
}

结果：

上面插入，下面删除。
第一个树高，第二个每次操作时间（毫秒）。