题目地址:
https://www.acwing.com/problem/content/255/
您需要写一种数据结构(可参考题目标题),来维护一些数,其中需要提供以下操作(以下排名都是指从小到大的排名。并列的数也要算进去):
1、插入数值
x
x
x。
2、删除数值
x
x
x(若有多个相同的数,应只删除一个)。
3、查询数值
x
x
x的排名(若有多个相同的数,应输出最小的排名)。
4、查询排名为
x
x
x的数值。
5、求数值
x
x
x的前驱(前驱定义为小于
x
x
x的最大的数)。
6、求数值
x
x
x的后继(后继定义为大于
x
x
x的最小的数)。
注意:数据保证查询的结果一定存在。
输入格式:
第一行为
n
n
n,表示操作的个数。接下来
n
n
n行每行有两个数opt和
x
x
x,opt表示操作的序号(opt范围是
1
∼
6
1\sim 6
1∼6)。
输出格式:
对于操作
3
,
4
,
5
,
6
3,4,5,6
3,4,5,6每行输出一个数,表示对应答案。
数据范围:
1
≤
n
≤
100000
1≤n≤100000
1≤n≤100000,所有数均在
−
1
0
7
−10^7
−107到
1
0
7
10^7
107内。
上面的 6 6 6个操作,很显然用BST来做比较适合。由于要保证插入和删除的时间复杂度维持在 O ( log n ) O(\log n) O(logn),我们需要对BST进行改进使之成为一棵平衡二叉搜索树。
法1:Treap(树堆)。Treap是一种二叉树结构,每个节点有两个权值,记为
k
k
k和
v
v
v,所有节点按照
k
k
k成为一个二叉搜索树,按照
v
v
v成为一个二叉堆(这个堆不一定是完全二叉树的形状,只是满足堆性质而已,即每个节点都大于等于两个孩子或小于等于两个孩子,对应最大堆和最小堆)。先介绍几个结论:
1、若任意两个节点的权值都不相同,Treap的结构是唯一的。这很容易证明。首先
v
v
v最大的节点应该是堆顶,那么其左子树就是由
k
k
k值小于这个节点
k
k
k值的节点组成,其余节点在右子树,接着递归建立左右子树即可,所以Treap的结构唯一。
2、对于
n
n
n个
k
k
k值确定的节点,假设它们
k
k
k值彼此不同,在随机产生
v
v
v的情况下,Treap的树高的期望是
O
(
log
n
)
O(\log n)
O(logn)的。所以其实
v
v
v这个权值是为了让BST平衡而设立的。这个结论说明了,如果在插入节点的时候,随机给其分配一个权值
v
v
v,那么对于这棵树查找、插入和删除的期望复杂度都是
O
(
log
n
)
O(\log n)
O(logn)。这是一个非常好的效果。而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
p
p的。但是实际写的时候最好都pushup以下,以保证旋转之后的size都正确。
设堆是最大堆(这其实无所谓,最小还是最大本质是一样的)。我们来考虑题目里的
6
6
6个操作都如何实现:
1、插入
x
x
x。先按照普通BST的方式递归插入
x
x
x,插入到叶子上后,回溯之前看一下被插入的那棵子树的树根是否违反了堆性质,如果违反了,则将其通过左旋或者右旋旋转到树根,接着回溯。最后记得pushup;
2、删除
x
x
x。首先有几个比较容易处理的情况,如果当前树空则不用删了,直接返回;如果当前节点的key大于
x
x
x,则去左子树里删;如果当前节点的key小于
x
x
x,则去右子树里删;否则说明当前节点的key等于
x
x
x,如果当前节点的计数大于
1
1
1,则直接计数减
1
1
1就行了;否则说明当前节点的计数是
1
1
1,那么必须删掉当前节点了。如果当前节点是叶子,那直接返回null节点就行了。接下来是最复杂的部分,即要删不是叶子的当前节点。Treap删除节点的方式是让要删除的节点旋转到叶子,然后再去删除叶子。所以我们考虑怎么把当前节点做旋转。如果该节点的右子树为空,或者右子树不空但是左孩子不空且其
v
v
v值大于右孩子的
v
v
v值,则可以右旋而不违反堆性质;否则的话说明右子树不空,并且或者左子树空,或者左孩子
v
v
v值小于右孩子
v
v
v值,此时做左旋。旋转完成之后去对应的子树进行递归删除。由于每次旋转都会使得要被删除的节点向下走,其总会成为叶子而被删除。最后记得pushup;
3、查询
x
x
x的排名。如果树空则返回
0
0
0,说明
x
x
x不存在。如果
x
x
x小于树根则去左子树取出其在左子树里的排名;如果
x
x
x大于树根,则其排名是其在右子树里的排名加上左子树的size再加上当前节点的计数;否则说明
x
x
x等于当前节点,其排名就是左子树的size加
1
1
1;
4、求排名为
x
x
x的数。如果左子树的size大于等于
x
x
x了,说明第
x
x
x名的数在左子树,则进左子树找排名第
x
x
x的数;如果左子树的节点个数
s
s
s加上树根的计数
c
c
c小于
x
x
x了,则说明第
x
x
x名的数在右子树里,则需要在右子树里去找排名第
x
−
s
−
c
x-s-c
x−s−c名的数;否则说明第
x
x
x的数就是当前树根,返回树根的
k
k
k值;
5、求小于
x
x
x的最大的数。如果树空则返回
−
∞
-\infty
−∞。如果
x
x
x小于等于树根,则去左子树里找;否则说明
x
x
x大于树根,那么小于
x
x
x的最大的数要么就是树根,要么去右子树里找,递归找到后取两者大的即可;
6、求大于
x
x
x的最小的数。与
5
5
5类似,如果树空则返回
+
∞
+\infty
+∞。如果
x
x
x大于等于树根,则去右子树里找;否则说明
x
x
x小于树根,那么大于
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(\log n) O(logn),空间 O ( n ) O(n) O(n), n 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
x
x进行分裂,可以分裂成两棵BST,左边那棵树的key都小于等于
x
x
x,右边那棵树的key都大于
x
x
x。可以递归地做。如果树空则分裂成两个空树。否则看一下树根与
x
x
x的关系,如果树根小于等于
x
x
x,那么树根及其左子树不需要分裂,只需要分裂右子树就行了,分裂完之后将分裂出来的小于等于
x
x
x的树连到原树根的右儿子处,大于
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 ) O(\log n) O(logn)的时间内合并两棵树。合并主要是为了维护堆性质。例如我们合并的两棵BBST是 x x x和 y y y,那么如果 x x x的val大于 y y y的val,我们就知道 y y y需要挂在 x x x的右子树,于是可以递归地将 x x x右子树与 y y y合并,作为新的右子树, x x x是新树树根;否则如果 x x x的val小于等于 y y 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
6
6个操作就好做了:
1、插入
x
x
x。先将原树按照
x
x
x分裂成
p
,
r
p,r
p,r,然后将
p
p
p按照
x
−
1
x-1
x−1分裂成
p
,
q
p,q
p,q,这样
p
<
x
,
q
=
x
,
r
>
x
p<x,q=x,r>x
p<x,q=x,r>x,看一下
q
q
q是否为空,如果是,则new出key为
x
x
x的节点,然后合并
p
,
q
,
r
p,q,r
p,q,r三棵树;否则将
q
q
q的cnt加
1
1
1,再合并
p
,
q
,
r
p,q,r
p,q,r;
2、删除
x
x
x。和操作
1
1
1一样按
x
x
x分裂成
p
,
q
,
r
p,q,r
p,q,r,看一下
q
q
q是否为空,如果空则直接将
p
,
q
,
r
p,q,r
p,q,r合并作为新树根;如果
q
q
q不空,则将
q
q
q的cnt减
1
1
1,接着看
q
q
q的cnt是否是
0
0
0,如果是,则合并
p
,
r
p,r
p,r作为新树根;
3、查询
x
x
x的排名。可以与普通Treap一样做,也可以按照
x
−
1
x-1
x−1分裂,第一棵树的size加
1
1
1就是
x
x
x的排名,算出来之后再合并回来。
4、求排名为
x
x
x的数。和普通Treap一样。
5、求小于
x
x
x的最大的数。按
x
−
1
x-1
x−1分裂,找第一棵树的最大值即可。找到之后合并回来。
6、求大于
x
x
x的最小的数。按
x
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
A
A,一开始全是
0
0
0,如果要插入一个数
x
x
x,则让
A
[
x
]
A[x]
A[x]加
1
1
1,如果是删除则让
A
[
x
]
A[x]
A[x]减去
1
1
1。这样
A
[
i
]
A[i]
A[i]代表
i
i
i存在的个数。由于题目中数字可能有负数,需要做离散化。若干查询操作实现如下:
1、查询
x
x
x的排名。求
A
A
A的前
x
−
1
x-1
x−1个数的前缀和,再加上
1
1
1就是
x
x
x的排名。
2、查询排名为
k
k
k的是哪个数。二分出最小的
x
x
x使得
∑
A
[
1
:
x
]
≥
k
\sum A[1:x]\ge k
∑A[1:x]≥k,这个
x
x
x即为所求。
3、求
x
x
x的前驱。求
k
=
∑
A
[
1
:
x
−
1
]
k=\sum A[1:x-1]
k=∑A[1:x−1],然后再求排名第
k
k
k的数是几即可。
4、求
x
x
x的后继。求
k
=
∑
A
[
1
:
x
]
k=\sum A[1:x]
k=∑A[1:x],然后求排名第
k
+
1
k+1
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;
}
预处理时间 O ( n log n ) O(n\log n) O(nlogn),操作 1 , 2 , 3 1,2,3 1,2,3时间 O ( log n ) O(\log n) O(logn),操作 4 , 5 , 6 4,5,6 4,5,6时间 O ( log 2 n ) O(\log^2 n) O(log2n),空间 O ( n ) O(n) O(n)。