洛谷P3369 AVL树

对洛谷题解上的代码进行了修改，删除了一些不必要的更新操作，以及修改了个人认为原来的代码不合理之处（可能是并不能理解原来的代码为何这么写） ，并添加了自己的注释

#include <bits/stdc++.h>
using namespace std;
struct AVLnode {
    // 数据
    int data;
    // 节点高度
    int high;
    // 相同数据的个数
    int freq;
    // 树的大小
    int size;
    // 左子节点
    AVLnode *ls;
    // 右子节点
    AVLnode *rs;
    // 默认构造函数，初始化数据，高度，相同数据个数，树的大小，左子节点，右子节点
    AVLnode() : data(0), high(1), freq(1), size(1), ls(NULL), rs(NULL) {}
    // 带参构造函数，初始化数据a，高度，相同数据个数，树的大小，左子节点，右子节点
    AVLnode(int a) : data(a), high(1), freq(1), size(1), ls(NULL), rs(NULL) {}
};
// 获取节点的大小
int GetSize(AVLnode* p) {
    if (p == NULL)
        return 0;
    return p->size;
}
// 获取节点的高度
int GetHigh(AVLnode* p) {
    if (p == NULL)
        return 0;
    return p->high;
}
// 更新节点的大小、高度
void update(AVLnode*& p) {
    // 该节点的大小为左子树的大小、右子树的大小和该节点相同数据的个数之和
    p->size = GetSize(p->ls) + GetSize(p->rs) + p->freq;
    // 该节点的高度为左子树的高度 + 右子树的高度 + 自身贡献高度1
    p->high = max(GetHigh(p->ls), GetHigh(p->rs)) + 1;
}
// 对树的旋转、插入、删除操作都有可能改变根节点，需要传引用
// 对树的查询不改变树的结构，不需要传引用

// 左左不平衡，将该节点右旋
void LeftLeft(AVLnode*& p) {
    AVLnode* q = p->ls;
    p->ls = q->rs;
    q->rs = p;
    update(p);
    update(q);
    p = q;
}
// 右右不平衡，将该节点左旋
void RightRight(AVLnode*& p) {
    AVLnode* q = p->rs;
    p->rs = q->ls;
    q->ls = p;
    update(p);
    update(q);
    p = q;
}
// 左右不平衡，将该节点的左子节点左旋，转化为左左不平衡情形，再将该节点右旋
void LeftRight(AVLnode*& p) {
    RightRight(p->ls);
    LeftLeft(p);
}
// 右左不平衡，将该节点的右子节点右旋，转化为右右不平衡情形，再将该节点左旋
void RightLeft(AVLnode*& p) {
    LeftLeft(p->rs);
    RightRight(p);
}
// 插入数据x
void Insert(AVLnode*& p, int x) {
    // 若根节点为空，说明这个子树为空树，只需要创建一个新节点即可
    if (p == NULL) { 
        p = new AVLnode(x);
        return;
    }
    // 若插入的数据x，在树中存在，只需要将相同数据的个数+1，并更新该节点的高度和相同数据个数
    if (p->data == x) {
        ++p->freq;
        update(p);
        return;
    }
    // 若插入的数据x小于目前节点p的数据，则往目前节点的左子树插入数据x，并更新节点p的高度
    // 接着调节树使得树仍处于平衡状态
    // 由于往目前节点左子树插入数据x，所以只可能导致树“左倾”(?)，故只会发生左右不平衡或左左不平衡
    // 在发生树“左倾”的情况下时
    // 当插入的数据x小于左节点数据时，发生左左不平衡
    // 当插入的数据x大于（不可能等于，若等于，则往左子树插入数据x时不会导致树不平衡）右节点数据时，发生左右不平衡
    if (x < p->data) {
        Insert(p->ls, x), update(p);
        if (GetHigh(p->ls) - GetHigh(p->rs) == 2) {
            if (x < p->ls->data)
                LeftLeft(p);
            else
                LeftRight(p);
        }
    }
    // 若插入的数据x大于（不会等于，因为等于的情况已经在上面return掉了）目前节点p的数据，则往目前节点的右子树插入数据x，并更新节点p的高度
    // 调整平衡的方式与上面类似
    else {
        Insert(p->rs, x), update(p);
        if (GetHigh(p->rs) - GetHigh(p->ls) == 2) {
            if (x > p->rs->data)
                RightRight(p);
            else
                RightLeft(p);
        }
    }
}
// 删除数据x
void Erase(AVLnode*& p, int x) {
    // 该节点为空，什么也不用做
    if (p == NULL)
        return;
    // 若删除的数据x比目前节点数据小，在左子树里删，比目前节点大，在右子树里删
    // 同插入，需要调整平衡
    // 若删左子树的数据，则可能发生“右倾”，反之，发生“左倾”
    // 不妨考虑“右倾”情形，此时无法通过删除的数据x来判断是右右不平衡还是右左不平衡。
    // 只能通过右节点的右子树高度和右节点的左子树高度比大小来判断是何种情形。
    // 需要注意右节点的左子树高度和右节点的右子树高度在发生“右倾”的情况下也是有可能相等的，这和插入情形不同。
    // 此时调整平衡的方式也应该是右右不平衡的方式（即只需要左旋）。
    // “左倾”情形同理。
    if (x < p->data) {
        Erase(p->ls, x), update(p);
        if (GetHigh(p->rs) - GetHigh(p->ls) == 2) {
            if (GetHigh(p->rs->rs) >= GetHigh(p->rs->ls))
                RightRight(p);
            else
                RightLeft(p);
        }
    } else if (x > p->data) {
        Erase(p->rs, x), update(p);
        if (GetHigh(p->ls) - GetHigh(p->rs) == 2) {
            if (GetHigh(p->ls->ls) >= GetHigh(p->ls->rs))
                LeftLeft(p);
            else
                LeftRight(p);
        }
    }
    // 若删除的数据x和该节点的数据相等，则要删该节点
    // 若该节点相同数据的个数大于1，只需要把相同节点个数-1
    else {
        if (p->freq > 1) {
            --p->freq;
            update(p);
            return;
        }
        // 若左子树和右子树都存在，则找该节点的后继，即大于它的最小值，只需一开始找右子树，之后一路向左
        // 用该节点后继的节点代替该节点，删掉该节点，并更新
        // 调整平衡，只可能是“左倾”，调整方法同上
        if (p->ls && p->rs) {
            AVLnode* q = p->rs;
            while (q->ls)
                q = q->ls;
            p->freq = q->freq;
            p->data = q->data;
            q->freq = 1;
            Erase(p->rs, q->data);
            update(p);
            if (GetHigh(p->ls) - GetHigh(p->rs) == 2) {
                if (GetHigh(p->ls->ls) >= GetHigh(p->ls->rs))
                    LeftLeft(p);
                else
                    LeftRight(p);
            }
        }
        // 若只存在左子树或只存在右子树，子承父业，删掉父节点即可，又由于删之前是平衡的树，所以删后不会发生不平衡的情形
        // 若左子树和右子树都不存在，直接删即可，又由于删之前是平衡的树，所以删后不会发生不平衡的情形
        else {
            AVLnode* q = p;
            if (p->ls)
                p = p->ls;
            else if (p->rs)
                p = p->rs;
            else
                p = NULL;
            delete q;
            q = NULL;
        }
    }
}
// 查询val的排名
int GetRank(AVLnode* p, int val) {
    // 此种情况，结果为左子树大小+1
    if (p->data == val)
        return GetSize(p->ls) + 1;
    // 此种情况，找左子树中val的排名
    if (val < p->data)
        return GetRank(p->ls, val);
    // 此种情况，找右子树中val的排名，加上左子树的大小，加节点本身的个数
    return GetRank(p->rs, val) + GetSize(p->ls) + p->freq;
}
// 查询排名为rank的数据
int GetVal(AVLnode* p, int rank) {
    // 此种情况，结果为左子树中找排名rank的数据
    if (GetSize(p->ls) >= rank)
        return GetVal(p->ls, rank);
    // 此种情况，结果为节点本身的数据
    if (GetSize(p->ls) + p->freq >= rank)
        return p->data;
    // 此种情况，结果为右子树中找排名为rank去掉左子树大小再去掉节点本身个数的数据
    return GetVal(p->rs, rank - GetSize(p->ls) - p->freq);
}
// 查找val前驱
int GetPrev(AVLnode* p, int val) {
    // 最大值大于等于-10e7
    int ans = -(10e7 + 5); 
    AVLnode* q = p;
    while (q) {
        // 存在val这个节点，先找左子树，再一路向右，答案即为一路向右
        if (q->data == val) {
            if (q->ls) {
                q = q->ls;
                while (q->rs)
                    q = q->rs;
                ans = q->data;
            }
            break;
        }
        // 不存在val这个节点，就找小于val的最大值
        if (q->data < val && q->data > ans)
            ans = q->data;
        // 往下遍历
        q = q->data < val ? q->rs : q->ls;
    }
    return ans;
}
// 查找val后继
int GetNext(AVLnode* p, int val) {
    // 方法同上
    int ans = 10e7 + 5;
    AVLnode *q = p;
    while (q) {
        if (q->data == val) {
            if (q->rs) {
                q = q->rs;
                while (q->ls)
                    q = q->ls;
                ans = q->data;
            }
            break;
        }
        if (q->data > val && q->data < ans)
            ans = q->data;
        q = q->data < val ? q->rs : q->ls;
    }
    return ans;
}
int main() {
    int n;
    scanf("%d", &n);
    AVLnode* root = NULL;
    while (n--) {
        int opt, x;
        scanf("%d%d", &opt, &x);
        switch (opt) {
            case 1:
                Insert(root, x);
                break;
            case 2:
                Erase(root, x);
                break;
            case 3:
                printf("%d\n", GetRank(root, x));
                break;
            case 4:
                printf("%d\n", GetVal(root, x));
                break;
            case 5:
                printf("%d\n", GetPrev(root, x));
                break;
            case 6:
                printf("%d\n", GetNext(root, x));
                break;
        }
    }
    return 0;
}