【学习】Splay平衡树

#include <cstdio>
#include <string>
#include <cstring>
#include <cstdlib>
#include <iostream> 
#define fi first
#define se second
#define in(x) getint(&x)
#define pii pair<int,int>
#define FOR(i,m,n) for(int i=m;i<=n;i++)
using namespace std;
typedef long long ll;
const int maxn = 100010;
const int inf = 1e9;
const int mod = 1e9+7;
template <typename T> inline void getint(T* _num){
    (*_num) = 0; int _op = 1; bool _ok = 0; char _c;
    while(1){_c = getchar(); if(_c == '-') _op = -1;if(_c <= 57 && _c >= 48)
        { _ok = 1; (*_num) = (*_num)*10+_c-48;}else if(_ok){(*_num)*=_op;return;}}}
template <typename T> void outint(T _out)
{if(_out >= 10) outint(_out/10); putchar(_out%10+48);}
/*
***********此下splay代码来源于胜利一中GTY学长***********
    笔记整理--splay平衡树基础操作 
    时间：2016年7月16日 
    注释+整理by: yhf_2015 -- TADYZ 
*/
struct splay_node{
    int value, size, cnt;
    splay_node *pre, *ch[2];    // pre记录父亲   ch[]记录儿子 
    void update(){  // update()函数每次更改size的大小 
        size = ch[0]->size + ch[1]->size + cnt;
    }
    int get_wh(){   // get_wh()判断方向函数 
    // 0表示左子树，1表示右子树 
        return pre->ch[0] == this ? 0 : 1; 
    }
    void set_ch(int wh, splay_node *child);
} pool[maxn], *root, *null;
// pool[maxn]为内存池   root记录根   null为空的假节点 
// set_ch修改儿子函数，把child改为wh方向上的儿子 
void splay_node :: set_ch(int wh, splay_node *child){
    ch[wh] = child; // 把now的wh方向的儿子变为child 
    if(child != null) child->pre = this; 
    // 如果child不是尾节点（空节点）则让child的前驱赋值为now 
    update(); 
    // 易错注意：一定要记得上传！！！！
    // 易错注意：一定要记得上传！！！！
    // 易错注意：一定要记得上传！！！！ 
}
int top = 0;
inline splay_node *get_new(int value){
    // 从内存池里提取一个内存 
    splay_node *now = pool + (++top);
// *******赋初值 *******************************
    now->value = value;
    now->size = 1;
    now->cnt = 1;
    now->pre = now->ch[0] = now->ch[1] = null;
//***********************************************
    // 返回此变量 
    return now;
}
// **************节点的旋转********************** 
// 此处now是在下面的节点 
inline void rotate(splay_node *now){ 
    splay_node *old_father = now->pre, *grand = now->pre->pre;
    // old_father 记录父亲   grand 记录父亲的父亲 

    // *********设now在father的左边 *****************
    // 因为now的右节点一定比now要大，一定比now的father要小，否则不会转向左边
    // 所以说把now的右子树给其father，now自己的右指针指向father，把自己提升上去 
    // **********************************************

    int wh = now->get_wh(); // 得到now是父节点的左儿子还是右儿子 
    old_father->set_ch(wh, now->ch[wh^1]); // 让now节点的父节点拥有now的右子树 
    now->set_ch(wh^1, old_father); // 让now的右子树指针指向now原来的父亲 old_father 
    now->pre = grand; // 让now的父节点等于旋转前的祖父节点
    // 相当于把now向上提升，取代其父节点的位置 
    // 对应的动画演示在2016_RZ夏令营DAY4课件——平衡树.pptx 18页 
    if(grand != null) // grand 为空不用执行，此时now已经是根节点了，无需更新祖父节点 
    // 让grand原来指向old_father的指针重新指向转完后代替old_father的now
    // 否则grand还是指向old_father
    // *******注意：原来grand哪个方向指向old_father ，现在还是那个方向指向now *****
        grand->ch[grand->ch[0] == old_father ? 0:1] = now; 
    else root = now;
}

// **********************splay平衡树伸展 *************************

// 伸展为splay的核心操作
// 把常访问的节点移动向根，通过递归旋转实现 
inline void splay(splay_node *now, splay_node *tar){ 
// tar为要把now旋转到的地方的父节点（if(now->pre == tar) 停止旋转操作） 
// 每次splay操作有执行两次旋转，第一次旋转分情况，第二次旋转now，将其提升 
// 每次在for中调用旋转——代码优美^_^ 
    for(; now->pre != tar; rotate(now))
        if(now->pre->pre != tar) // 因为每次操作会影响到now的父节点，防止有奇数步翻转的情况 
        // 因为上面的for还有一次旋转操作，所以这里不再旋转 
        // 此处分为两种情况： 
        // 1.他们在一条直链上（在自己父节点的方向相同），此时先旋转now的father，再转now 
        // 2.他们不再一条直链上，此时直接旋转两次now 
        // 注意：第二次旋转在上面的for里面 
            now->get_wh() == now->pre->get_wh() ? rotate(now->pre) : rotate(now);
    // 如果tar == null意思是直接把now旋转到根节点，所以说直接让根节点的指针指向now 
    if(tar == null) root = now;
} 

// **********************splay 插入节点*******************************
// value为插入值得权值 
void insert(int value){
    splay_node *last = null, *now = root; // 两个指针，last指向空，now从root开始查询 
    splay_node *newone = get_new(value); // 申请新的权值为value的值
    // 这里可以维护队列，每次删除操作，把内存池编号加入队列，优先从队列里取内存，优化空间 
    while(now != null){ // 只要没搜到底，继续往下搜 
        last = now;
        // 因为是可重复的splay，如果比较时两值相同，now的个数增加，节点数量增加 
        if(newone->value == now->value){ 
            now->cnt ++, now->size ++; 
            splay(now, null);
            return;
        }
        // 如果新加入的比now要小，向左搜索，大向右搜索 
        if(newone->value < now->value) now = now->ch[0];
        else now = now->ch[1];
    }
    // 处理原来没有点的情况，这时候直接让now作为root 
    if(last == null) { root = newone; return; }
    // 和搜索到的最后一个点比较权值的大小，决定插入左边还是右边 
    else{
        if(newone->value < last->value) last->set_ch(0, newone);
        else last->set_ch(1, newone);
        splay(newone, null);    // 每次伸展，把新的点旋转到根 
    } 
}
// 查找元素是否存在，存在自动splay伸展到根 
inline splay_node *find(int value){
    splay_node *now = root; // 从根开始搜索 
    while(now != null){
        if(now->value == value) break;// 相等说明已经找到，退出返回 
        if(value < now->value) now = now->ch[0];
        else now = now->ch[1];
    }
    // 排除没有点的情况，防止RE 
    if(now != null) splay(now, null);
    return now;
} 
// **********************找前驱的函数 **************************
// 前驱是整个平衡树中第一个比value小的数 
// 先找到以value为根子树，第一次向左走，因为一定比他小，第二次有右子树就向右 
inline int pre(int value){
    int ans = -inf;
    // 不用find函数，直接从根开始找，其实都一样 
    splay_node *now = root;
    while(now != null)
        if(now->value < value){ // 小于value，更新ans，向右走 
            ans = max(ans, now->value);
            now = now->ch[1]; 
        }else now = now->ch[0]; // 大于向左走 
    return ans;
}
// ********************找后续的函数 ***************************
// 与找前驱的相同 
inline int next(int value){
    int ans = inf;
    splay_node *now = root;
    while(now != null)
        if(now->value > value){
            ans = min(ans, now->value);
            now = now->ch[0];
        }else now = now->ch[1];
    return ans;
}
// ********************删除节点操作 *******************************
void del(int value){
    // 删除value值的节点
    // 主要操作：先查找value的值，value被伸展到了根节点的位置
    splay_node *now = find(value);
    if(now == null) return; // 如果没有该值，直接返回 
    // 如果now的数量大于1，直接删除其中的一个，返回 
    if(now->cnt > 1) { now->cnt--; now->size--; return; }
    /*
    首先找到要删除的节点,然后把它伸展(splay)到根的位置。再分三种情况:
    1. 无左右儿子,直接删除节点。
    2. 只有一个儿子,令独生子变成根,删去该点。
    3. 左右儿子都有。找到左子树中权值最大的点,将其 splay 到左儿子的位置。
    然后将 整棵右子树挂在左儿子的右边。删除节点。 
    */
    // 处理删光的情况，根直接等于空 
    if(now->ch[0] == null && now->ch[1] == null) root = null;
    else if(now->ch[1] == null){
        now->ch[0]->pre = null;
        root = now->ch[0];
    }else if(now->ch[0] == null){
        now->ch[1]->pre = null;
        root = now->ch[1];
    }else{
        splay_node *tmp = now->ch[0];
        // 找左子树中权值最大的点 
        while(tmp->ch[1] != null) tmp = tmp->ch[1];
        splay(tmp, now);
        tmp->set_ch(1, now->ch[1]);
        tmp->pre = null;
        root = tmp;
    }
}
// 找到value的排名 -> 使用大小二分的方法
inline int get_rank(int value){
    splay_node *now = root;
    int left_size = 0;
    while(now != null){
        if(value == now->value){
            int ans = left_size + now->ch[0]->size + 1;
            splay(now, null);
            return ans;
        }
        else if(value < now->value) now = now->ch[0];
        else left_size += now->ch[0]->size + now->cnt, now = now->ch[1];
    }
    return -1;
}
// 找到平衡树中第k大的元素 
inline int kth(int k){
    splay_node *now = root;
    int left_size = 0;
    while(now != null){
        int tmp = left_size + now->ch[0]->size;
        if(tmp + 1 <= k && k <= tmp + now->cnt){
            splay(now, null);
            return now->value;
        }
        if(k <= tmp) now = now->ch[0];
        else left_size = tmp + now->cnt, now = now->ch[1];
    }
    return -1;
}
int main(){
    null = pool;
    null->value = 0;
    null->size = 0;
    null->cnt = 0;
    null->pre = null->ch[0] = null->ch[1] = null;
    root = null;
    int q; in(q);
    while(q--){
        int order, val; 
        in(order), in(val);
        switch(order){
            case 1: insert(val); break;
            case 2: del(val); break;
            case 3: printf("%d\n", get_rank(val)); break;
            case 4: printf("%d\n", kth(val)); break;
            case 5: printf("%d\n", pre(val)); break;
            case 6: printf("%d\n", next(val)); break;
        }
    }
    return 0;
}