【ACWing】253. 普通平衡树

记录算法题解

已于 2023-10-12 14:04:31 修改

阅读量460

点赞数 1

分类专栏： AC 数据结构文章标签：算法数据结构 c++

于 2021-07-23 06:45:19 首次发布

本文链接：https://blog.csdn.net/qq_46105170/article/details/118997891

版权

AC 数据结构专栏收录该内容

199 篇文章 5 订阅

订阅专栏

题目地址：

https://www.acwing.com/problem/content/255/

您需要写一种数据结构（可参考题目标题），来维护一些数，其中需要提供以下操作（以下排名都是指从小到大的排名。并列的数也要算进去）：
1、插入数值 $x$ 。
2、删除数值 $x$ （若有多个相同的数，应只删除一个）。
3、查询数值 $x$ 的排名（若有多个相同的数，应输出最小的排名）。
4、查询排名为 $x$ 的数值。
5、求数值 $x$ 的前驱(前驱定义为小于 $x$ 的最大的数)。
6、求数值 $x$ 的后继(后继定义为大于 $x$ 的最小的数)。
注意：数据保证查询的结果一定存在。

输入格式：
第一行为 $n$ ，表示操作的个数。接下来 $n$ 行每行有两个数opt和 $x$ ，opt表示操作的序号（opt范围是 $1\sim 6$ ）。

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

数据范围：
$1 \leq n \leq 100000$ ，所有数均在 $10^7$ 到 $10^7$ 内。

上面的 $6$ 个操作，很显然用BST来做比较适合。由于要保证插入和删除的时间复杂度维持在 $O(\log n)$ ，我们需要对BST进行改进使之成为一棵平衡二叉搜索树。

法1：Treap（树堆）。Treap是一种二叉树结构，每个节点有两个权值，记为 $k$ 和 $v$ ，所有节点按照 $k$ 成为一个二叉搜索树，按照 $v$ 成为一个二叉堆（这个堆不一定是完全二叉树的形状，只是满足堆性质而已，即每个节点都大于等于两个孩子或小于等于两个孩子，对应最大堆和最小堆）。先介绍几个结论：
1、若任意两个节点的权值都不相同，Treap的结构是唯一的。这很容易证明。首先 $v$ 最大的节点应该是堆顶，那么其左子树就是由 $k$ 值小于这个节点 $k$ 值的节点组成，其余节点在右子树，接着递归建立左右子树即可，所以Treap的结构唯一。
2、对于 $n$ 个 $k$ 值确定的节点，假设它们 $k$ 值彼此不同，在随机产生 $v$ 的情况下，Treap的树高的期望是 $O(\log n)$ 的。所以其实 $v$ 这个权值是为了让BST平衡而设立的。这个结论说明了，如果在插入节点的时候，随机给其分配一个权值 $v$ ，那么对于这棵树查找、插入和删除的期望复杂度都是 $O(\log n)$ 。这是一个非常好的效果。而Treap的维护比别的BBST比如红黑树、AVL树都要简单很多，只需要做两种旋转操作就可以了，即左旋和右旋。

左旋和右旋的目的，是将不满足堆性质的节点旋转上去或者旋转下来。先介绍左旋和右旋如何实现。首先来看右旋，如图所示：
在这里插入图片描述
右旋的输入是树根节点，右旋操作就是要让该节点的左孩子顺时针向上旋转成为新的树根，而原树根成为新树根的右孩子。那么原树根的左孩子的右孩子就没地方去了，由于它是大于新树根的，所以要将它往右子树里放，而它又是小于原树根的，所以就将其放到新树根右孩子的左边作为其左子树即可。具体可以看上图。右旋又叫做zig，代码如下：

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

这里的pushup操作和线段树里的是一模一样的意思，都是以孩子节点的信息更新自己。
对于左旋有类似操作，如下图：
在这里插入图片描述
左旋又叫做zag，代码如下：

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

非动图如下：
在这里插入图片描述
注意，在旋转里的pushup(p)一般是不需要写的，因为在insert和remove函数中最后一定会更新 $p$ 的。但是实际写的时候最好都pushup以下，以保证旋转之后的size都正确。

设堆是最大堆（这其实无所谓，最小还是最大本质是一样的）。我们来考虑题目里的 $6$ 个操作都如何实现：
1、插入 $x$ 。先按照普通BST的方式递归插入 $x$ ，插入到叶子上后，回溯之前看一下被插入的那棵子树的树根是否违反了堆性质，如果违反了，则将其通过左旋或者右旋旋转到树根，接着回溯。最后记得pushup；
2、删除 $x$ 。首先有几个比较容易处理的情况，如果当前树空则不用删了，直接返回；如果当前节点的key大于 $x$ ，则去左子树里删；如果当前节点的key小于 $x$ ，则去右子树里删；否则说明当前节点的key等于 $x$ ，如果当前节点的计数大于 $1$ ，则直接计数减 $1$ 就行了；否则说明当前节点的计数是 $1$ ，那么必须删掉当前节点了。如果当前节点是叶子，那直接返回null节点就行了。接下来是最复杂的部分，即要删不是叶子的当前节点。Treap删除节点的方式是让要删除的节点旋转到叶子，然后再去删除叶子。所以我们考虑怎么把当前节点做旋转。如果该节点的右子树为空，或者右子树不空但是左孩子不空且其 $v$ 值大于右孩子的 $v$ 值，则可以右旋而不违反堆性质；否则的话说明右子树不空，并且或者左子树空，或者左孩子 $v$ 值小于右孩子 $v$ 值，此时做左旋。旋转完成之后去对应的子树进行递归删除。由于每次旋转都会使得要被删除的节点向下走，其总会成为叶子而被删除。最后记得pushup；
3、查询 $x$ 的排名。如果树空则返回 $0$ ，说明 $x$ 不存在。如果 $x$ 小于树根则去左子树取出其在左子树里的排名；如果 $x$ 大于树根，则其排名是其在右子树里的排名加上左子树的size再加上当前节点的计数；否则说明 $x$ 等于当前节点，其排名就是左子树的size加 $1$ ；
4、求排名为 $x$ 的数。如果左子树的size大于等于 $x$ 了，说明第 $x$ 名的数在左子树，则进左子树找排名第 $x$ 的数；如果左子树的节点个数 $s$ 加上树根的计数 $c$ 小于 $x$ 了，则说明第 $x$ 名的数在右子树里，则需要在右子树里去找排名第 $x - s - c$ 名的数；否则说明第 $x$ 的数就是当前树根，返回树根的 $k$ 值；
5、求小于 $x$ 的最大的数。如果树空则返回 $-\infty$ 。如果 $x$ 小于等于树根，则去左子树里找；否则说明 $x$ 大于树根，那么小于 $x$ 的最大的数要么就是树根，要么去右子树里找，递归找到后取两者大的即可；
6、求大于 $x$ 的最小的数。与 $5$ 类似，如果树空则返回 $+\infty$ 。如果 $x$ 大于等于树根，则去右子树里找；否则说明 $x$ 小于树根，那么大于 $x$ 的最小的数要么就是树根，要么去左子树里找，递归找到后取两者小的即可；

代码如下：

#include <iostream>
#include <cstring>
using namespace std;

const int N = 100010, INF = 1e8;
int n;
// 注意tr[0]空着不用，视为null节点，其cnt和size都是0
struct Node {
	// 左右孩子的下标
    int l, r;
    // key是BST的权值，val是最大堆的权值
    int key, val;
    // cnt是本key有多少个，size表示以本node为树根的子树的节点总个数
    int cnt, size;
} tr[N];
int root, idx;

void pushup(int p) {
    tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}
// 右旋
void zig(int &p) {
    int q = tr[p].l;
    tr[p].l = tr[q].r, tr[q].r = p, p = q;
    pushup(tr[p].r);
}
// 左旋
void zag(int &p) {
    int q = tr[p].r;
    tr[p].r = tr[q].l, tr[q].l = p, p = q;
    pushup(tr[p].l);
}

// 创建新节点并返回其下标
int get_node(int key) {
    tr[++idx].key = key;
    tr[idx].val = rand();
    tr[idx].cnt = tr[idx].size = 1;
    return idx;
}

// 注意要传引用，因为我们要真实修改p为新的Treap的树根
void insert(int &p, int key) {
	// 如果树空，则new出新节点作为树根
    if (!p) p = get_node(key);
    else if (tr[p].key > key) {
    	// 如果key小于树根，则要向左子树里插入
        insert(tr[p].l, key);
        // 左子树完成插入后，看一下左子树的val是否与当前节点
        // 的val满足堆的关系，如果不满足，则将左子树右旋上来
        if (tr[tr[p].l].val > tr[p].val) zig(p);
    } else if (tr[p].key < key) {
    	// 如果key大于当前树根，则要向右子树插入
        insert(tr[p].r, key);
        // 右子树完成插入后，看一下右子树的val是否与当前节点
        // 的val满足堆的关系，如果不满足，则将右子树左旋上来
        if (tr[tr[p].r].val > tr[p].val) zag(p);
        // 否则说明当前节点的key就等于输入key，则计数加1
    } else tr[p].cnt++;
	
	// 插入完成后还需要更新当前节点
    pushup(p);
}

void remove(int &p, int key) {
	// 如果树空，那就不用删了，直接返回
    if (!p) return;
    // 如果key大于树根，则去右子树里删
    if (tr[p].key > key) remove(tr[p].l, key);
    // 如果key小于树根，则去左子树里删
    else if (tr[p].key < key) remove(tr[p].r, key);
    // 否则说明key等于树根
    else {
    	// 如果key之前的次数大于1，那么不用删节点，只需要删cnt
        if (tr[p].cnt > 1) tr[p].cnt--;
        // 否则说明要删当前节点了
        else if (tr[p].l || tr[p].r) {
        	// 如果右子树空，或者左子树不空且v值大于右孩子v值，则需要右旋，否则左旋
            if (!tr[p].r || (!tr[p].l && tr[tr[p].l].val > tr[tr[p].r].val)) {
                zig(p);
                remove(tr[p].r, key);
            } else {
                zag(p);
                remove(tr[p].l, key);
            }
            // 如果当前节点是叶子，那直接将其变成空节点null就行了，null的下标是0
        } else p = 0;
    }

	// 删除完成后还需要更新当前节点
    pushup(p);
}

int get_rank_by_key(int p, int key) {
    if (!p) return 0;
    if (tr[p].key > key) return get_rank_by_key(tr[p].l, key);
    else if (tr[p].key < key) return tr[tr[p].l].size + tr[p].cnt + get_rank_by_key(tr[p].r, key);
    else return tr[tr[p].l].size + 1;
}

int get_key_by_rank(int p, int rank) {
    if (!p) return INF;
    if (tr[tr[p].l].size >= rank) return get_key_by_rank(tr[p].l, rank);
    else if (tr[tr[p].l].size + tr[p].cnt < rank) return get_key_by_rank(tr[p].r, rank - tr[tr[p].l].size - tr[p].cnt);
    else return tr[p].key;
}

int get_prev(int p, int key) {
    if (!p) return -INF;
    if (tr[p].key >= key) return get_prev(tr[p].l, key);
    return max(tr[p].key, get_prev(tr[p].r, key));
}

int get_next(int p, int key) {
    if (!p) return INF;
    if (tr[p].key <= key) return get_next(tr[p].r, key);
    return min(tr[p].key, get_next(tr[p].l, key));
}

int main() {
    scanf("%d", &n);
    while (n--) {
        int opt, x;
        scanf("%d%d", &opt, &x);
        if (opt == 1) insert(root, x);
        else if (opt == 2) remove(root, x);
        else if (opt == 3) printf("%d\n", get_rank_by_key(root, x));
        else if (opt == 4) printf("%d\n", get_key_by_rank(root, x));
        else if (opt == 5) printf("%d\n", get_prev(root, x));
        else if (opt == 6) printf("%d\n", get_next(root, x));
    }

    return 0;
}

每个操作期望时间复杂度都是 $O(\log n)$ ，空间 $O (n)$ ， $n$ 是当前Treap里节点个数。

也可以采取简易写法：

#include <iostream>
using namespace std;

const int N = 100010, INF = 1e9;
int root, idx;
struct Node {
  // ch[0]代表左孩子，ch[1]代表右孩子
  int ch[2], key, val, cnt, sz;
} tr[N];

#define lc(x) tr[x].ch[0]
#define rc(x) tr[x].ch[1]

int get_node(int key) {
  tr[++idx] = {{0, 0}, key, rand(), 1, 1};
  return idx;
}

#define pushup(p) tr[p].sz = tr[lc(p)].sz + tr[rc(p)].sz + tr[p].cnt

// d = 0代表左旋，d = 1代表右旋
void rotate(int &p, int d) {
  int q = tr[p].ch[d ^ 1];
  tr[p].ch[d ^ 1] = tr[q].ch[d], tr[q].ch[d] = p, p = q;
  pushup(tr[p].ch[d]), pushup(p);
}

void insert(int &p, int key) {
  if (!p) p = get_node(key);
  else if (tr[p].key == key) tr[p].cnt++;
  else {
    int d = tr[p].key < key;
    insert(tr[p].ch[d], key);
    if (tr[tr[p].ch[d]].val > tr[p].val) rotate(p, d ^ 1);
  }
  pushup(p);
}

void remove(int &p, int key) {
  if (!p) return;
  if (tr[p].key == key) {
    if (tr[p].cnt > 1) tr[p].cnt--;
    // 如果左右孩子都非空，则左旋右旋都可以，取决于堆性质
    else if (lc(p) && rc(p)) {
      // 如果我们要保持大根堆，左孩子大就要右旋，反之亦然
      int d = tr[lc(p)].val > tr[rc(p)].val;
      rotate(p, d);
      remove(tr[p].ch[d], key);
      // 如果两个子树里有一个为空，那p直接取非空的那个就行
    } else p = lc(p) ?: rc(p);
  } else remove(tr[p].ch[tr[p].key < key], key);
  pushup(p);
}

int get_rank_by_key(int p, int key) {
  if (!p) return 0;
  if (tr[p].key > key) return get_rank_by_key(lc(p), key);
  else if (tr[p].key < key)
    return tr[lc(p)].sz + tr[p].cnt + get_rank_by_key(rc(p), key);
  return tr[lc(p)].sz + 1;
}

int get_key_by_rank(int p, int rank) {
  if (!p) return INF;
  if (tr[lc(p)].sz >= rank) return get_key_by_rank(lc(p), rank);
  if (tr[lc(p)].sz + tr[p].cnt < rank)
    return get_key_by_rank(rc(p), rank - tr[lc(p)].sz - tr[p].cnt);
  return tr[p].key;
}

int get_prev(int p, int key) {
  if (!p) return -INF;
  if (tr[p].key >= key) return get_prev(lc(p), key);
  return max(tr[p].key, get_prev(rc(p), key));
}

int get_next(int p, int key) {
  if (!p) return INF;
  if (tr[p].key <= key) return get_next(rc(p), key);
  return min(tr[p].key, get_next(lc(p), key));
}

int main() {
  int n;
  scanf("%d", &n);
  while (n--) {
    int opt, x;
    scanf("%d%d", &opt, &x);
    if (opt == 1) insert(root, x);
    else if (opt == 2) remove(root, x);
    else if (opt == 3) printf("%d\n", get_rank_by_key(root, x));
    else if (opt == 4) printf("%d\n", get_key_by_rank(root, x));
    else if (opt == 5) printf("%d\n", get_prev(root, x));
    else printf("%d\n", get_next(root, x));
  }
}

有时候需要让每个节点只表示一个数出现了一次，即cnt总为0。代码如下：

#include <iostream>
using namespace std;

const int N = 100010, INF = 1e8;
int root, idx;
struct Node {
  int ch[2], key, val, sz;
} tr[N];

#define lc(x) tr[x].ch[0]
#define rc(x) tr[x].ch[1]

int get_node(int key) {
  tr[++idx] = {{0, 0}, key, rand(), 1};
  return idx;
}

#define pushup(p) tr[p].sz = tr[lc(p)].sz + tr[rc(p)].sz + 1

void rotate(int& p, int d) {
  int q = tr[p].ch[d ^ 1];
  tr[p].ch[d ^ 1] = tr[q].ch[d], tr[q].ch[d] = p, p = q;
  pushup(tr[p].ch[d]), pushup(p);
}

void insert(int& p, int key) {
  if (!p) {
    p = get_node(key);
    return;
  }

  int d = tr[p].key < key;
  insert(tr[p].ch[d], key);
  if (tr[tr[p].ch[d]].val > tr[p].val) rotate(p, d ^ 1);
  pushup(p);
}

void remove(int& p, int key) {
  if (!p) return;
  if (tr[p].key == key) {
    if (lc(p) && rc(p)) {
      int d = tr[lc(p)].val > tr[rc(p)].val;
      rotate(p, d);
      remove(tr[p].ch[d], key);
    } else p = lc(p) ?: rc(p);
  } else remove(tr[p].ch[tr[p].key < key], key);
  // 尤其要注意，如果p = 0，是不能做pushup的，pushup完了
  // 空节点的sz就变为1了，整个就都错了
  if (p) pushup(p);
}

int get_rank_by_key(int p, int key) {
  if (!p) return INF;
  if (key <= tr[p].key)
    return min(1 + tr[lc(p)].sz, get_rank_by_key(lc(p), key));
  else return 1 + tr[lc(p)].sz + get_rank_by_key(rc(p), key);
}

int get_key_by_rank(int p, int rank) {
  if (rank <= tr[lc(p)].sz) return get_key_by_rank(lc(p), rank);
  else if (rank > tr[lc(p)].sz + 1)
    return get_key_by_rank(rc(p), rank - 1 - tr[lc(p)].sz);
  return tr[p].key;
}

int get_prev(int p, int key) {
  if (!p) return -INF;
  if (key <= tr[p].key) return get_prev(lc(p), key);
  return max(tr[p].key, get_prev(rc(p), key));
}

int get_next(int p, int key) {
  if (!p) return INF;
  if (tr[p].key <= key) return get_next(rc(p), key);
  return min(tr[p].key, get_next(lc(p), key));
}

int main() {
  int n;
  scanf("%d", &n);
  while (n--) {
    int opt, x;
    scanf("%d%d", &opt, &x);
    if (opt == 1) insert(root, x);
    else if (opt == 2) remove(root, x);
    else if (opt == 3) printf("%d\n", get_rank_by_key(root, x));
    else if (opt == 4) printf("%d\n", get_key_by_rank(root, x));
    else if (opt == 5) printf("%d\n", get_prev(root, x));
    else printf("%d\n", get_next(root, x));
  }
}

法2：FHQ Treap（范浩强Treap）。其不是利用旋转，而是利用分裂合并来实现上面的所有操作的。分裂操作参考https://blog.csdn.net/qq_46105170/article/details/119298420。可以参考下图：
在这里插入图片描述
在图中，灰色是原来的树，其按照某个值 $x$ 进行分裂，可以分裂成两棵BST，左边那棵树的key都小于等于 $x$ ，右边那棵树的key都大于 $x$ 。可以递归地做。如果树空则分裂成两个空树。否则看一下树根与 $x$ 的关系，如果树根小于等于 $x$ ，那么树根及其左子树不需要分裂，只需要分裂右子树就行了，分裂完之后将分裂出来的小于等于 $x$ 的树连到原树根的右儿子处，大于 $x$ 的树就是第二棵树；如果树根大于 $x$ ，那么类似，树根及其右子树不需要分裂，只分裂左子树即可。代码如下：

// 将子树p按照key来分裂，x存小于等于key的树，y存大于key的树
void split(int p, int key, int &x, int &y) {
    if (!p) x = y = 0;
    else {
        if (tr[p].key <= key) {
            x = p;
            split(tr[p].r, key, tr[p].r, y);
        } else {
            y = p;
            split(tr[p].l, key, x, tr[p].l);
        }
        // 分裂完了之后还需要pushup一下
		pushup(p);
    }
}

合并操作是分裂的逆操作，即给定两棵BBST，第一棵的每个节点都小于第二棵树，那么可以在 $O(\log n)$ 的时间内合并两棵树。合并主要是为了维护堆性质。例如我们合并的两棵BBST是 $x$ 和 $y$ ，那么如果 $x$ 的val大于 $y$ 的val，我们就知道 $y$ 需要挂在 $x$ 的右子树，于是可以递归地将 $x$ 右子树与 $y$ 合并，作为新的右子树， $x$ 是新树树根；否则如果 $x$ 的val小于等于 $y$ 的val，可以类似操作。代码如下：

int merge(int x, int y) {
    if (!x || !y) return x + y;
    if (tr[x].val > tr[y].val) {
        tr[x].r = merge(tr[x].r, y);
        pushup(x);
        return x;
    } else {
        tr[y].l = merge(x, tr[y].l);
        pushup(y);
        return y;
    }
}

有了这两个操作，题中的 $6$ 个操作就好做了：
1、插入 $x$ 。先将原树按照 $x$ 分裂成 $p, r$ ，然后将 $p$ 按照 $x - 1$ 分裂成 $p, q$ ，这样 $p < x, q = x, r > x$ ，看一下 $q$ 是否为空，如果是，则new出key为 $x$ 的节点，然后合并 $p, q, r$ 三棵树；否则将 $q$ 的cnt加 $1$ ，再合并 $p, q, r$ ；
2、删除 $x$ 。和操作 $1$ 一样按 $x$ 分裂成 $p, q, r$ ，看一下 $q$ 是否为空，如果空则直接将 $p, q, r$ 合并作为新树根；如果 $q$ 不空，则将 $q$ 的cnt减 $1$ ，接着看 $q$ 的cnt是否是 $0$ ，如果是，则合并 $p, r$ 作为新树根；
3、查询 $x$ 的排名。可以与普通Treap一样做，也可以按照 $x - 1$ 分裂，第一棵树的size加 $1$ 就是 $x$ 的排名，算出来之后再合并回来。
4、求排名为 $x$ 的数。和普通Treap一样。
5、求小于 $x$ 的最大的数。按 $x - 1$ 分裂，找第一棵树的最大值即可。找到之后合并回来。
6、求大于 $x$ 的最小的数。按 $x$ 分裂，找第二棵树的最小值即可，找到之后合并回来。

代码如下：

#include <iostream>
using namespace std;

const int N = 1e5 + 10, INF = 1e8;
int n;

struct Node {
    int l, r;
    int key, val;
    int cnt, size;
} tr[N];
int root, idx;
int x, y, z;

int get_node(int key) {
    tr[++idx].key = key;
    tr[idx].val = rand();
    tr[idx].size = tr[idx].cnt = 1;
    return idx;
}

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

// 将树p按照key分裂成两棵树，第一棵小于等于key，树根是x，第二棵大于key，树根是y
void split(int p, int key, int &x, int &y) {
	// 如果树空，则分裂为两棵空树
    if (!p) x = y = 0;
    else {
        if (tr[p].key <= key) {
        	// 如果当前树根的key小于等于key，那么第一棵树树根就是p，
        	// 继续分裂右子树，将右子树分裂出的第一棵树接到p的右儿子上
            x = p;
            split(tr[p].r, key, tr[p].r, y);
        } else {
            y = p;
            split(tr[p].l, key, x, tr[p].l);
        }
        // 分裂完之后更新p的信息
	    pushup(p);
    }
}

// 将两棵BST合并为一棵
int merge(int x, int y) {
	// 如果一棵为空，则直接返回另一棵
    if (!x || !y) return x ^ y;
    // 如果x的val大于y的val，那么y应该接到x的右子树里，
    // 则合并x的右子树与y，接到x的右边，新树根就是x
    if (tr[x].val > tr[y].val) {
        tr[x].r = merge(tr[x].r, y);
        pushup(x);
        return x;
    } else {
        tr[y].l = merge(x, tr[y].l);
        pushup(y);
        return y;
    }
}

void insert(int key) {
    split(root, key, x, z);
    split(x, key - 1, x, y);
    if (!y) y = get_node(key);
    else {
        tr[y].cnt++;
        tr[y].size++;
    }

    root = merge(merge(x, y), z);
}

void remove(int key) {
    split(root, key, x, z);
    split(x, key - 1, x, y);
    if (y) {
        tr[y].cnt--;
        tr[y].size--;
    }

    if (!tr[y].cnt) y = 0;
    root = merge(merge(x, y), z);
}

int get_rank_by_key(int key) {
    split(root, key - 1, x, y);
    int rank = tr[x].size + 1;
    root = merge(x, y);
    return rank;
}

int get_key_by_rank(int rank) {
    int p = root;
    while (p) {
        if (tr[tr[p].l].size >= rank)
            p = tr[p].l;
        else if (tr[tr[p].l].size + tr[p].cnt < rank) {
            rank -= tr[tr[p].l].size + tr[p].cnt;
            p = tr[p].r;
        } else return tr[p].key;
    }

    return INF;
}

int get_prev(int key) {
    split(root, key - 1, x, y);
    int p = x;
    while (tr[p].r) p = tr[p].r;
    root = merge(x, y);
    return tr[p].key;
}

int get_next(int key) {
    split(root, key, x, y);
    int p = y;
    while (tr[p].l) p = tr[p].l;
    root = merge(x, y);
    return tr[p].key;
}

int main() {
    scanf("%d", &n);
    while (n--) {
        int opt, a;
        scanf("%d%d", &opt, &a);
        if (opt == 1) insert(a);
        else if (opt == 2) remove(a);
        else if (opt == 3) printf("%d\n", get_rank_by_key(a));
        else if (opt == 4) printf("%d\n", get_key_by_rank(a));
        else if (opt == 5) printf("%d\n", get_prev(a));
        else if (opt == 6) printf("%d\n", get_next(a));
    }

    return 0;
}

所有操作时空复杂度一样。

法3：树状数组。开一个数组 $A$ ，一开始全是 $0$ ，如果要插入一个数 $x$ ，则让 $A [x]$ 加 $1$ ，如果是删除则让 $A [x]$ 减去 $1$ 。这样 $A [i]$ 代表 $i$ 存在的个数。由于题目中数字可能有负数，需要做离散化。若干查询操作实现如下：
1、查询 $x$ 的排名。求 $A$ 的前 $x - 1$ 个数的前缀和，再加上 $1$ 就是 $x$ 的排名。
2、查询排名为 $k$ 的是哪个数。二分出最小的 $x$ 使得 $\sum A[1:x]\ge k$ ，这个 $x$ 即为所求。
3、求 $x$ 的前驱。求 $k=\sum A[1:x-1]$ ，然后再求排名第 $k$ 的数是几即可。
4、求 $x$ 的后继。求 $k=\sum A[1:x]$ ，然后求排名第 $k + 1$ 的数是几即可。

代码如下：

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

const int N = 1e5 + 10;
int n;
int tr[N], op[N], v[N], a[N];
int sz, idx;

int lowbit(int x) {
    return x & -x;
}

void add(int k, int x) {
    for (int i = k; i <= n; i += lowbit(i)) tr[i] += x;
}

int sum(int k) {
    int res = 0;
    for (int i = k; i; i -= lowbit(i)) res += tr[i];
    return res;
}

// 求x离散化之后的值是多少
int get(int x) {
    return lower_bound(a + 1, a + 1 + idx, x) - a;
}

void insert(int x) {
    add(x, 1);
}

void remove(int x) {
    add(x, -1);
}

int get_rank_by_key(int x) {
    return sum(x - 1) + 1;
}

int get_key_by_rank(int rk) {
    int l = 1, r = n;
    while (l < r) {
        int mid = l + r >> 1;
        if (sum(mid) >= rk) r = mid;
        else l = mid + 1;
    }
	// 这里的l是离散化之后的答案，还需要还原回来才是正确答案
    return a[l];
}

int get_prev(int x) {
    return get_key_by_rank(sum(x - 1));
}

int get_next(int x) {
    return get_key_by_rank(sum(x) + 1);
}

int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%d%d", &op[i], &v[i]);
        if (op[i] != 4) a[++idx] = v[i];
    }

	// 排序完去重
    sort(a + 1, a + 1 + idx);
    sz = unique(a + 1, a + 1 + idx) - (a + 1);

    for (int i = 1; i <= n; i++) {
        int opt = op[i], x = v[i];
        // 如果不是第4个操作，则求一下x离散化之后的值
        if (opt != 4) x = get(v[i]);

        if (opt == 1) insert(x);
        else if (opt == 2) remove(x);
        else if (opt == 3) printf("%d\n", get_rank_by_key(x));
        else if (opt == 4) printf("%d\n", get_key_by_rank(x));
        else if (opt == 5) printf("%d\n", get_prev(x));
        else if (opt == 6) printf("%d\n", get_next(x));
    }

    return 0;
}