splay的写法

Posted 2020-12-10 sdlang

距离上次写splay已经过去了10个月了，今天高兴地重拾了平衡树，赶紧过来写一下自己的写法，以后好养成习惯

需要解释的尽量在代码里注释了，就不过多说了

检查x是父亲的左儿子还是右儿子

int chk(int x)
{
    return ch[fa[x]][1] == x;
}

pushup

void pushup(int x)
{
    size[x] = size[ch[x][0]] + size[ch[x][1]] + cnt[x];
}

旋转

void rot(int x)
{
    int y = fa[x],z = fa[y],k = chk(x);   //y是x的父亲,z是x的爷爷
    ch[z][chk(y)] = x;          //把x转到z的儿子替代y
    fa[x] = z;
    ch[y][k] = ch[x][k ^ 1];    //y继承x对着的儿子
    fa[ch[x][k ^ 1]] = y;
    ch[x][k ^ 1] = y;        //y变成x的对着的儿子
    fa[y] = x;
    pushup(y);
    pushup(x);              //记着从下往上pushup
    //这个应该是比较短而且清晰的旋转了吧QAQ
}

双旋

void splay(int x,int goal)    //goal为0是转到根
{
    while (fa[x] != goal)
    {
        int y = fa[x],z = fa[y];
        if (z != goal)
        {
            if (chk(x) == chk(y))    //‘三点共线‘
                rot(y);
            else
                rot(x);
        }
        rot(x);
    }
    if (!goal)
        rt = x;
}

插入

void insert(int x)
{
    int u = rt,f = 0;        //f是父亲
    while (u && val[u] != x)
    {
        f = u;
        u = ch[u][x > val[u]];   //根据大小关系判断往左走还是往右走
    }
    if (u)        //这个位置存在数
        cnt[u]++;
    else
    {
        u = ++node_cnt;    //新建节点
        fa[u] = f;
        size[u] = cnt[u] = 1;
        val[u] = x;
        ch[u][1] = ch[u][0] = 0;
        if (f)
            ch[f][x > val[f]] = u;
    }
    splay(u,0);        //splay上去更新
}

把x所在位置转到根节点

void find(int x)
{
    int u = rt;
    if (!u)    //树莫得了
        return;
    while (ch[u][x > val[u]] && x != val[u])  //能往下走而且没走到
        u = ch[u][x > val[u]];
    splay(u,0);
}

查前驱或后继

int nxt(int x,int p)    //0是前驱 1是后继
{
    find(x);      //先把要找的数转到根节点
    int u = rt;
    if (val[u] > x && p)   //根节点的数比一下
        return u;
    if (val[u] < x && !p)
        return u;
    u = ch[u][p];
    while (ch[u][p ^ 1])
        u = ch[u][p ^ 1];
    //最后这三行是因为前驱是在根节点左子树的最大值,后继在根节点右子树的最小值
    //所以先走到左右子树,然后反着跳
    return u;
}

删除

void del(int x)
{
    int l = nxt(x,0),r = nxt(x,1);   //前驱后继
    splay(l,0);
    splay(r,l);
    int &d = ch[r][0];
    //我们把前驱转到根节点,把后继转到前驱的儿子
    //那么后继肯定是前驱的右儿子,x是后继的左儿子且一定为叶子节点
    if (cnt[d] > 1)
    {
        cnt[d]--;
        splay(d,0);
    }
    else
        d = 0;
}

第k大

int kth(int x)
{
    int u = rt;
    if (size[u] < x)    //就没有这么多数
        return 0;
    while (114514)
    {
        int l = ch[u][0];
        if (x > size[l] + cnt[u])    //比左儿子节点数和当前数的个数多
        {
            x -= size[l] + cnt[u];
            u = ch[u][1];            //往右儿子走
        }
        else
            if (size[l] >= x)   //左儿子有那么多数
                u = l;          //往左儿子走
            else
                return val[u];  //不然就是这里了
    }
}