平衡二叉树-AVL c/c++代码实现

参考:https://www.bilibili.com/video/BV1rt411j7Ff?t=703大佬视频
一份代码,代码中有注释,对应着洛谷的P3369 【模板】普通平衡树

/*Keep on going Never give up*/
#pragma GCC optimize(3,"Ofast","inline")
#include <bits/stdc++.h>
const int maxn = 2e5+10;
const int MaxN = 0x3f3f3f3f;
const int MinN = 0xc0c0c00c;
typedef long long ll;
const int mod = 100000000;
using namespace std;
struct Node{
    int l,r,val,height,size;
}a[maxn];
int cnt,root;
inline void creat_node(int &now,int val){  //新创建一个结点
    now=++cnt;
    a[now].val=val;
    a[now].size=1;
}
inline int get_height(int now){   //计算树的高度差,左-右
    return a[a[now].l].height-a[a[now].r].height;
}
inline void update(int now){       //更新树的高度差以及元素个数
    a[now].size=a[a[now].l].size+a[a[now].r].size+1;
    a[now].height=max(a[a[now].l].height,a[a[now].r].height)+1;
}
inline void lrotate(int &now){   //左旋
    int r=a[now].r;     //把他右子树的结点暂存一下,一会需要以右孩子为根节点
    a[now].r=a[a[now].r].l;   //把根节点右子树的左子树挂在原来根节点的上右子树上
    a[r].l=now;  //把原来的根节点挂在原来的右子树(现在的根节点)上去
    now=r;   //现在的根节点更新
    update(a[now].l),update(now);  //检查树是否需要更新
}

inline void rrotate(int &now){  //右旋
    int l=a[now].l;
    a[now].l=a[a[now].l].r;
    a[l].r=now;
    now=l;
    update(a[now].r),update(now);
}

inline void check(int &now){  //检查树是否需要旋转
    int nh=get_height(now);
    if(nh>1){    //如果左边-右边高度大于1
        int lh=get_height(a[now].l);
        if(lh>0) rrotate(now);   //LL情况
        else lrotate(a[now].l),rrotate(now);  //LR情况
    }
    else if(nh<-1){
        int rh=get_height(a[now].r);
        if(rh<0) lrotate(now);   //RR情况
        else rrotate(a[now].r),lrotate(now);       //RL情况
    }
    else if(now) update(now);
}
inline void insert(int &now,int val){  //插入
    if(!now) creat_node(now,val);   //如果没有结点,插入
    else if(val<a[now].val) insert(a[now].l,val);   //搜索树的性质,小了往左差,大了往右插
    else insert(a[now].r,val);
    check(now);
}

inline int ifind(int &now,int fa){  //查找后继点
    int ret;
    if(!a[now].l){
        ret=now;
        a[fa].l=a[now].r;
    }
    else{
        ret=ifind(a[now].l,now);
        check(now);
    }
    return ret;
}

inline void del(int &now,int val){   //删稠某点
    if(val==a[now].val){
        int l=a[now].l,r=a[now].r;
        if(!l||!r) now=l+r;  //两儿子都没有或者是只有一个儿子
        else{
            now=ifind(r,r);  //找后继
            if(now!=r)  a[now].r=r; //如果后继是他本身
            a[now].l=l;
        }
    }
    else if(val<a[now].val) del(a[now].l,val);
    else del(a[now].r,val);
    check(now);
}

inline int get_rank(int val){  //找出某个数的名次
    int now=root,rank=1;
    while (now){
        if(val<=a[now].val) now=a[now].l;  //往树的左孩子走
        else{
            rank+=a[a[now].l].size+1;   //往树的右孩子走,同时减去加上他小的个数(他的左子树的数都比他小,所以他的总排名肯定要比左边的所有数都高,所以先用rank加上左数大小)
            now=a[now].r;   //更新目前的根节点
        }
    }
    return rank;
}
inline int get_num(int rank){  //找出名次的数字
    int now=root;
    while (now){
        if(a[a[now].l].size+1==rank) break;   //如果找到了直接break
        else if(a[a[now].l].size>=rank) now=a[now].l;  //同理
        else{
            rank-=a[a[now].l].size+1;   //往树的右孩子走,同时减去比他小的个数(他的左子树的数都比他小,所以直接减去左子树的大小)
            now=a[now].r;
        }
    }
    return a[now].val;
}

inline int read(){
    int s=0,w=1;
    char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')w=-1;ch=getchar();}
    while(ch>='0'&&ch<='9') s=s*10+ch-'0',ch=getchar();
    return s*w;
}

int main()
{
    int n;
    cin>>n;
    for(int i=0;i<n;i++){
        int ch,x;
        ch=read(),x=read();
        if(ch==1) insert(root,x);
        if(ch==2) del(root,x);
        if(ch==3) printf("%d\n",get_rank(x));
        if(ch==4) printf("%d\n",get_num(x));
        if(ch==5) printf("%d\n",get_num(get_rank(x)-1));
        if(ch==6) printf("%d\n",get_num(get_rank(x+1)));
    }
    return 0;
}

数据结构算法学习 文章被收录于专栏

算法学习记录

全部评论

相关推荐

评论
点赞
收藏
分享
牛客网
牛客企业服务