C7 普通平衡树 • EH's Blog

题目情境#

题目链接：https://www.luogu.com.cn/problem/P3369 ↗

题目描述

您需要动态地维护一个可重集合 $M$ ，并且提供以下操作：

向 $M$ 中插入一个数 $x$ 。
从 $M$ 中删除一个数 $x$ （若有多个相同的数，应只删除一个）。
查询 $M$ 中有多少个数比 $x$ 小，并且将得到的答案加一。
查询如果将 $M$ 从小到大排列后，排名位于第 $x$ 位的数。
查询 $M$ 中 $x$ 的前驱（前驱定义为小于 $x$ ，且最大的数）。
查询 $M$ 中 $x$ 的后继（后继定义为大于 $x$ ，且最小的数）。

对于操作 3,5,6，不保证当前可重集中存在数 $x$ 。

输入格式

第一行为 $n$ ，表示操作的个数,下面 $n$ 行每行有两个数 $\text{opt}$ 和 $x$ ， $\text{opt}$ 表示操作的序号（ $1 \leq \text{opt} \leq 6$ ）

输出格式

对于操作 $3,4,5,6$ 每行输出一个数，表示对应答案。

输入 #1

plaintext

输出 #1

106465
84185
492737

plaintext

【数据范围】对于 $100\%$ 的数据， $1\le n \le 10^5$ ， $|x| \le 10^7$

算法解析#

过难，略

Python代码实现#

import sys

class Node:
    def __init__(self):
        self.ch = [0, 0]  # 左右孩子
        self.fa = 0  # 父节点
        self.v = 0  # 节点值
        self.cnt = 0  # 值出现次数
        self.siz = 0  # 子树大小

    def init(self, p, v1):
        self.fa = p
        self.v = v1
        self.cnt = self.siz = 1


def ls(x):
    return tr[x].ch[0]


def rs(x):
    return tr[x].ch[1]


def pushup(x):
    tr[x].siz = tr[ls(x)].siz + tr[rs(x)].siz + tr[x].cnt


def rotate(x):
    y = tr[x].fa
    z = tr[y].fa
    k = 1 if tr[y].ch[1] == x else 0
    tr[z].ch[1 if tr[z].ch[1] == y else 0] = x
    tr[x].fa = z
    tr[y].ch[k] = tr[x].ch[k ^ 1]
    tr[tr[x].ch[k ^ 1]].fa = y
    tr[x].ch[k ^ 1] = y
    tr[y].fa = x
    pushup(y)
    pushup(x)


def splay(x, k):
    while tr[x].fa != k:
        y = tr[x].fa
        z = tr[y].fa
        if z != k:
            if (ls(y) == x) ^ (ls(z) == y):
                rotate(x)
            else:
                rotate(y)
        rotate(x)
    if k == 0:
        global root
        root = x


def insert(v):
    global root, tot
    x = root
    p = 0
    while x != 0 and tr[x].v != v:
        p = x
        x = tr[x].ch[1 if v > tr[x].v else 0]
    if x != 0:
        tr[x].cnt += 1
    else:
        tot += 1
        x = tot
        tr[p].ch[1 if v > tr[p].v else 0] = x
        tr[x].init(p, v)
    splay(x, 0)


def find(v):
    global root
    x = root
    while tr[x].ch[1 if v > tr[x].v else 0] != 0 and v != tr[x].v:
        x = tr[x].ch[1 if v > tr[x].v else 0]
    splay(x, 0)


def getpre(v):
    global root
    find(v)
    x = root
    if tr[x].v < v:
        return x
    x = ls(x)
    while rs(x) != 0:
        x = rs(x)
    splay(x, 0)
    return x


def getsuc(v):
    global root
    find(v)
    x = root
    if tr[x].v > v:
        return x
    x = rs(x)
    while ls(x) != 0:
        x = ls(x)
    splay(x, 0)
    return x


def del_(v):
    pre = getpre(v)
    suc = getsuc(v)
    splay(pre, 0)
    splay(suc, pre)
    del_node = tr[suc].ch[0]
    if tr[del_node].cnt > 1:
        tr[del_node].cnt -= 1
        splay(del_node, 0)
    else:
        tr[suc].ch[0] = 0
        splay(suc, 0)


def getrank(v):
    insert(v)
    res = tr[ls(root)].siz
    del_(v)
    return res


def getval(k):
    global root
    x = root
    while True:
        if k <= tr[ls(x)].siz:
            x = ls(x)
        elif k <= tr[ls(x)].siz + tr[x].cnt:
            break
        else:
            k -= tr[ls(x)].siz + tr[x].cnt
            x = rs(x)
    splay(x, 0)
    return tr[x].v


N = 110001
INF = (1 << 30) + 1
tr = [Node() for _ in range(N)]
root = 0
tot = 0


insert(-INF)
insert(INF)
n = int(sys.stdin.readline())
for _ in range(n):
    op, x = map(int, sys.stdin.readline().split())
    if op == 1:
        insert(x)
    elif op == 2:
        del_(x)
    elif op == 3:
        print(getrank(x))
    elif op == 4:
        print(getval(x + 1))
    elif op == 5:
        print(tr[getpre(x)].v)
    else:
        print(tr[getsuc(x)].v)

python

C++代码实现#

#include <iostream>
using namespace std;

#define ls(x) tr[x].ch[0]
#define rs(x) tr[x].ch[1]
const int N=1100010, INF=(1<<30)+1;
struct node{
  int ch[2]; //儿
  int fa; //父
  int v;  //点权
  int cnt; //点权次数
  int siz; //子树大小
  void init(int p,int v1){
    fa=p, v=v1;
    cnt=siz=1;
  }
}tr[N];
int root,tot; //根,节点个数

void pushup(int x){ //上传
  tr[x].siz=tr[ls(x)].siz+tr[rs(x)].siz+tr[x].cnt;
}
void rotate(int x){ //旋转
  int y=tr[x].fa, z=tr[y].fa, k=tr[y].ch[1]==x; //y的右儿是x
  tr[z].ch[tr[z].ch[1]==y]=x, tr[x].fa=z; //z的儿是x,x的父是z
  tr[y].ch[k]=tr[x].ch[k^1], tr[tr[x].ch[k^1]].fa=y; //y的儿是x的异儿,x的异儿的父是y
  tr[x].ch[k^1]=y, tr[y].fa=x; //x的异儿是y,y的父是x
  pushup(y), pushup(x); //自底向上push
}
void splay(int x, int k){ //伸展
  while(tr[x].fa!=k){ //折线转xx,直线转yx
    int y=tr[x].fa, z=tr[y].fa;
    if(z!=k) (ls(y)==x)^(ls(z)==y)?rotate(x):rotate(y);
    rotate(x);
  }
  if(!k) root=x; //k=0时,x转到根
}
void insert(int v){ //插入
  int x=root, p=0;
  //x走到空节点或走到目标点结束
  while(x&&tr[x].v!=v) p=x,x=tr[x].ch[v>tr[x].v];
  if(x) tr[x].cnt++; //目标点情况
  else{ //空节点情况
    x=++tot;
    tr[p].ch[v>tr[p].v]=x;
    tr[x].init(p,v);
  }
  splay(x, 0);
}
void find(int v){ //找到v并转到根
  int x=root;
  while(tr[x].ch[v>tr[x].v]&&v!=tr[x].v)
    x=tr[x].ch[v>tr[x].v];
  splay(x, 0);
}
int getpre(int v){ //前驱
  find(v);
  int x=root;
  if(tr[x].v<v) return x;
  x=ls(x);
  while(rs(x)) x=rs(x);
  splay(x, 0);
  return x;
}
int getsuc(int v){ //后继
  find(v);
  int x=root;
  if(tr[x].v>v) return x;
  x=rs(x);
  while(ls(x)) x=ls(x);
  splay(x, 0);
  return x;
}
void del(int v){ //删除
  int pre=getpre(v);
  int suc=getsuc(v);
  splay(pre,0), splay(suc,pre);
  int del=tr[suc].ch[0];
  if(tr[del].cnt>1)
    tr[del].cnt--, splay(del,0);
  else
    tr[suc].ch[0]=0, splay(suc,0);
}
int getrank(int v){ //排名
  insert(v);
  int res=tr[tr[root].ch[0]].siz;
  del(v);
  return res;
}
int getval(int k){ //数值
  int x=root;
  while(true){
    if(k<=tr[ls(x)].siz) x=ls(x);
    else if(k<=tr[ls(x)].siz+tr[x].cnt) break;
    else k-=tr[ls(x)].siz+tr[x].cnt, x=rs(x);
  }
  splay(x, 0);
  return tr[x].v;
}
int main(){
  insert(-INF);insert(INF); //哨兵
  int n,op,x; scanf("%d", &n);
  while(n--){
    scanf("%d%d", &op, &x);
    if(op==1) insert(x);
    else if(op==2) del(x);
    else if(op==3) printf("%d\n",getrank(x));
    else if(op==4) printf("%d\n",getval(x+1));
    else if(op==5) printf("%d\n",tr[getpre(x)].v);
    else printf("%d\n",tr[getsuc(x)].v);
  }
}

更好的写法：

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;

using pii = pair<int,int>;

tree<
    pii,
    null_type,
    less<pii>,
    rb_tree_tag,
    tree_order_statistics_node_update
> tr;

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n; cin >> n;
    int id = 0;
    const int INF = 1e9;

    while(n--){
        int op, x; cin >> op >> x;

        if(op == 1){                     // 插入 x
            tr.insert({x, ++id});
        }
        else if(op == 2){                // 删除一个 x
            auto it = tr.lower_bound({x, -INF});
            if(it != tr.end() && it->first == x) tr.erase(it);
        }
        else if(op == 3){                // x 的排名
            cout << tr.order_of_key({x, -INF}) + 1 << '\n';
        }
        else if(op == 4){                // 第 x 小
            cout << tr.find_by_order(x - 1)->first << '\n';
        }
        else if(op == 5){                // x 的前驱
            auto it = tr.lower_bound({x, -INF});
            --it;
            cout << it->first << '\n';
        }
        else if(op == 6){                // x 的后继
            auto it = tr.lower_bound({x, INF});
            cout << it->first << '\n';
        }
    }
}