平衡树

平衡树

Treap实现

思路:

利用堆的性质, 让二叉搜索数满足堆的性质,从而达到logn的高度.
模板
具体解释看注释,注释也不多(逃)

代码:

/*
 * 平衡数Treap模板
 * Treap 可以理解为一棵树加上一个堆, 通过对每个节点赋予一个随机值
 * 在满足堆的性质的同时满足二叉搜索树的性质, 保证树的高度尽量为logn,
 * 这样就不会出现较坏的情况
 */
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e6+10;
const int INF = 1e9;

struct Treap
{
    int l, r;//左儿子和右儿子的标号
    int data;//节点的随机值
    int val;//节点维护的值
    int cnt;//节点当前值的数量
    int size;//当前子树的值
}treap[MAXN];
int cnt, root, n;//节点总数

void PushUp(int p)
{
    //向上更新
    treap[p].size = treap[treap[p].l].size+treap[treap[p].r].size+treap[p].cnt;
}

void RotateRight(int &p)
{
    //右旋操作
    int q = treap[p].l;
    treap[p].l = treap[q].r;
    treap[q].r = p;
    p = q;
    PushUp(treap[p].r);
    PushUp(p);
}

void RotateLeft(int &p)
{
    //左旋操作
    int q = treap[p].r;
    treap[p].r = treap[q].l;
    treap[q].l = p;
    p = q;
    PushUp(treap[p].l);
    PushUp(p);
}

int NewNode(int val)
{
    ++cnt;
    treap[cnt].val = val;
    treap[cnt].cnt = treap[cnt].size = 1;
    treap[cnt].data = rand();
    treap[cnt].l = treap[cnt].r = -1;//表示没有子节点
    return cnt;
}

void Build()
{
    //初始化建个树用来判断边界条件
    NewNode(-INF), NewNode(INF);
    root = 1;//根节点初始化
    treap[root].r = 2;
    PushUp(root);
}

void Insert(int &p, int val)
{
    //插入节点
    if (p == -1)
    {
        //说明这个节点没有使用
        p = NewNode(val);
        return;
    }
    if (val == treap[p].val)
    {
        //查询到相同值
        treap[p].cnt++;
        PushUp(p);
        return;
    }
    if (val < treap[p].val)
    {
        //往左节点移动
        Insert(treap[p].l, val);
        //根据data值来调整树的高度
        if (treap[p].data < treap[treap[p].l].data)
            RotateRight(p);
    }
    else
    {
        //往右节点移动
        Insert(treap[p].r, val);
        if (treap[p].data < treap[treap[p].r].data)
            RotateLeft(p);
    }
    PushUp(p);//更新
}

void Remove(int &p, int val)
{
    //删除节点即将要删除的点给移动到叶子节点再进行删除即可.
    if (p == -1)
        return;
    //不存在这个值
    if (val == treap[p].val)
    {
        if (treap[p].cnt > 1)
        {
            treap[p].cnt--;
            PushUp(p);
            return;
        }
        if (treap[p].l != -1 || treap[p].r != -1)
        {
            if (treap[p].r == -1 || treap[treap[p].l].data > treap[treap[p].r].data)
            {
                RotateRight(p);
                Remove(treap[p].r, val);
            }
            else
            {
                RotateLeft(p);
                Remove(treap[p].l, val);
            }
            PushUp(p);
        }
        else
        {
            p = -1;
            return;
        }
    }
    if (val < treap[p].val)
        Remove(treap[p].l, val);
    else
        Remove(treap[p].r, val);
    PushUp(p);
}

int GetRankByVal(int p, int val)
{
    //得到某个值的排名
    if (p == -1)
        return 0;
    if (val == treap[p].val)
        return treap[treap[p].l].size + 1;
    if (val < treap[p].val)
        return GetRankByVal(treap[p].l, val);
    return GetRankByVal(treap[p].r, val)+treap[treap[p].l].size+treap[p].cnt;
}

int GetValByRank(int p, int rank)
{
    //找到排名对应的值
    if (p == -1)
        return INF;
    if (treap[treap[p].l].size >= rank)
        return GetValByRank(treap[p].l, rank);
    if (treap[treap[p].l].size + treap[p].cnt >= rank)
        return treap[p].val;
    return GetValByRank(treap[p].r, rank-treap[treap[p].l].size-treap[p].cnt);
}

int GetPre(int val)
{
    int ans = 1;//treap[1].val = -INF
    int p = root;
    while (p != -1)
    {
        if (val == treap[p].val)
        {
            if (treap[p].l != -1)
            {
                p = treap[p].l;
                while (treap[p].r != -1)//在左子树上一直往右走
                    p = treap[p].r;
                ans = p;
            }
            break;
        }
        if (treap[p].val < val && treap[p].val > treap[ans].val)
            ans = p;
        if (val < treap[p].val)
            p = treap[p].l;
        else
            p = treap[p].r;
    }
    return treap[ans].val;
}

int GetNext(int val)
{
    int ans = 2;//treap[2].val = INF
    int p = root;
    while (p != -1)
    {
        if (treap[p].val == val)
        {
            if (treap[p].r != -1)
            {
                p = treap[p].r;
                while (treap[p].l != -1)
                    p = treap[p].l;
                ans = p;
            }
            break;
        }
        if (treap[p].val > val && treap[p].val < treap[ans].val)
            ans = p;
        if (val < treap[p].val)
            p = treap[p].l;
        else
            p = treap[p].r;
    }
    return treap[ans].val;
}

int main()
{
    Build();
    scanf("%d", &n);
    int opt, val;
    while (n--)
    {
        scanf("%d%d", &opt, &val);
        switch (opt)
        {
            case 1:
                Insert(root, val);
                break;
            case 2:
                Remove(root, val);
                break;
            case 3:
                printf("%d
", GetRankByVal(root, val)-1);
                break;
            case 4:
                printf("%d
", GetValByRank(root, val+1));
                break;
            case 5:
                printf("%d
", GetPre(val));
                break;
            case 6:
                printf("%d
", GetNext(val));
                break;
        }
    }

    return 0;
}