【数据结构】平衡树

实现功能：
插入数值
删除数值
查询某排名的数字
查询某数值的排名
查询前驱后继
cpp 复制代码
const int N = 100010, INF = 1e8;

int n;
struct Node
{
    int l, r; // 左右子结点编号
    int key, val; // key:结点本身的值 val:为了使二叉树平衡的随机数
    int cnt, size; // cnt:当前结点的数出现了多少次 size:以当前结点为根的子树中有多少结点
}tr[N];

int root, idx; // root:根结点编号 idx:计数器,记录结点编号用到哪了

void pushup(int p) // 更新p
{
    tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}

int getnode(int key) // 创建值为key的新结点
{
    tr[ ++ idx].key = key;
    tr[idx].cnt = 1;
    tr[idx].size = 1;
    tr[idx].val = rand();
    return idx;
}

void zig(int &p) // 右旋 (以p为根)
{
    int q = tr[p].l;
    tr[p].l = tr[q].r;
    tr[q].r = p;
    p = q;
    pushup(tr[p].r), pushup(p);
}

void zag(int &p) // 左旋 (以p为根)
{
    int q = tr[p].r;
    tr[p].r = tr[q].l;
    tr[q].l = p;
    p = q;
    pushup(tr[p].l), pushup(p);
}

void build() // 创建treap
{
    getnode(-INF), getnode(INF);
    root = 1, tr[1].r = 2;
    pushup(root);

    if (tr[1].val < tr[2].val) zag(root);
}

void insert(int &p, int key) // 当前在p结点，要创建值为key的结点
{
    if (!p) p = getnode(key); // p不存在
    else if (tr[p].key == key) tr[p].cnt ++ ; // 原本就有key的结点
    else if (tr[p].key > key) // 当前结点的key比新结点的key大
    {
        insert(tr[p].l, key);
        if (tr[tr[p].l].val > tr[p].val) zig(p);
    }
    else // 当前结点的key比新结点的key小
    {
        insert(tr[p].r, key);
        if (tr[tr[p].r].val > tr[p].val) zag(p);
    }
    pushup(p); // 更新当前结点
}

void remove(int &p, int key) // 当前在p结点，要删去值为key的结点
{
    if (!p) return;
    else if (tr[p].key == key) // 找到要删去的结点
    {
        if (tr[p].cnt > 1) tr[p].cnt -- ; // 想删去的key不止一个
        else if (tr[p].l || tr[p].r) // 左右子结点至少存在一个
        {
            if (!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val)
            {
                zig(p);
                remove(tr[p].r, key);
            }
            else if (!tr[p].l || tr[tr[p].l].val < tr[tr[p].r].val)
            {
                zag(p);
                remove(tr[p].l, key);
            }
        }
        else p = 0; // 叶子结点
    }
    else if (tr[p].key > key) remove(tr[p].l, key); // 当前结点值小于要删去的key
    else if (tr[p].key < key) remove(tr[p].r, key); // 当前结点值大于要删去的key

    pushup(p);
}

int getrank_bykey(int p, int key) // 当前结点为p，给定key找rank
{
    if (!p) return 0;
    else if (tr[p].key == key) return tr[tr[p].l].size + 1; // 当前结点即为要找的结点
    else if (tr[p].key > key) return getrank_bykey(tr[p].l, key); // 要找的结点在当前左子树中
    else if (tr[p].key < key) return tr[tr[p].l].size + tr[p].cnt + getrank_bykey(tr[p].r, key); // 要找的结点在当前右子树中
}

int getkey_byrank(int p, int rank) // 当前结点为p，给定rank找key
{
    if (!p) return INF;
    else if (tr[tr[p].l].size >= rank) return getkey_byrank(tr[p].l, rank); // 要找的结点在当前左子树
    else if (tr[tr[p].l].size + tr[p].cnt >= rank) return tr[p].key; // 要找的结点即为当前结点
    else return getkey_byrank(tr[p].r, rank - tr[tr[p].l].size - tr[p].cnt); // 要找的结点在当前右子树
}

int getprev(int p, int key) // 当前结点为p，找小于key的最大值
{
    if (!p) return -INF;
    else if (tr[p].key >= key) return getprev(tr[p].l, key); // 小于key的结点在当前结点的左子树
    else if (tr[p].key < key) return max(tr[p].key, getprev(tr[p].r, key)); // 小于key的结点在当前结点or右子树
}

int getnext(int p, int key) // 当前结点为p，找大于key的最小值
{
    if (!p) return INF;
    else if (tr[p].key <= key) return getnext(tr[p].r, key); // 大于key的结点在当前结点的右子树
    else if (tr[p].key > key) return min(tr[p].key, getnext(tr[p].l, key)); // 大于key的结点在当前结点or右子树
}