众所周知，平衡树是提高数据结构。

平衡树

平衡树首先是一颗二叉搜索树 (BST)。二叉搜索树的定义如下：

  1. 空树是二叉搜索树。
  2. 若二叉搜索树的左子树不为空，则其左子树上所有点的附加权值均小于其根节点的值。
  3. 若二叉搜索树的右子树不为空，则其右子树上所有点的附加权值均大于其根节点的值。
  4. 二叉搜索树的左右子树均为二叉搜索树。

我们不难发现，在二叉搜索树上可以很方便的进行如下操作：

  1. 从大到小遍历
  2. 查询最值
  3. 插入元素
  4. 删除元素
  5. 求元素排名
  6. 求排名为 k 的元素

除了第一个操作复杂度为 O(n)O(n)，其他这些操作的复杂度都是 O(h)O(h) 的，其中 hh 为树的高度。最坏情况下，二叉搜索树会退化成一条链，于是复杂度都退化为 O(n)O(n)，这是我们不希望看到的。因此，出现了平衡树的概念。平衡指：每一个结点的左子树和右子树高度差最多为 11，这样就可以保证所有复杂度都是在 O(log⁡n)O(\log n) 级别完成的。

set

首先，在你写平衡树之前，先想想 STL 的 set （或者 multiset） 够不够你用。

  1. 从大到小遍历：可以，复杂度 O(n)O(n)
  2. 查询最值：可以，复杂度 O(log⁡n)O(\log n)
  3. 插入元素：可以，复杂度 O(log⁡n)O(\log n)
  4. 删除元素：可以，复杂度 O(log⁡n)O(\log n)
  5. 求元素排名：*可以，但是复杂度 O(n)O(n)
  6. 求排名为 k 的元素：*可以，但是复杂度 O(n)O(n)

set 无法做到查询排名和查询排名为 k 元素的操作，只能一个一个遍历得到。但是 set 支持查询某个元素，并获得其的迭代器，复杂度也为 O(log⁡n)O(\log n) （用 lower_bound 和 upper_bound 成员函数）。就是说可以 O(log⁡n)O(\log n) 求出前驱，后继。并且可以 O(1)O(1) 遍历到下一个元素。

Treap

核心思想是给每个节点赋一个随机的优先级，然后将这个优先级按照堆的性质维护。

堆的性质：每个节点的键值都大于等于/小于等于其父亲的键值。（我们这里用小根堆）

节点结构

用数组模拟一棵树。

1
2
3
4
5
6
7
8
9
10
const int N = 1e5+5;
int l[N],r[N]; // 左右儿子
int val[N],pri[N]; // 值和随机的优先级
int cnt[N],sz[N]; // 重复次数,子树大小
int root; // 树根
int m; // 节点个数

void maintain(int i){ // 维护 sz 的大小
    sz[i] = cnt[i] + sz[l[i]] + sz[r[i]];
}

旋转操作

旋转操作可以在不破坏二叉搜索树性质的情况下，调整其结构。

如图，可以看到确实旋转后仍然满足二叉搜索树的性质：

  • 对于左旋操作

    1. 根节点的右节点的左节点 改为 根节点
    2. 根节点的右节点 改为 根节点右节点的左节点
    3. 根节点右节点 成为 新的根节点

  • 对于右旋操作

    1. 根节点的左节点的右节点 改为 根节点
    2. 根节点的左节点 改为 根节点左节点的右节点
    3. 根节点左节点 成为 新的根节点

对于这段代码的记忆，请画图，然后就很清晰了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void lrotate(int &k){ // 左旋
    int t=r[k]; // 这将是新根节点
    r[k]=l[t]; // 原根节点的右节点是新根节点的左节点
    l[t]=k; // 新根节点的左节点是原根节点
    sz[t]=sz[k]; // 更新新根节点的 sz
    maintain(k); // 更新原根节点的 sz
    k = t; // 引用，修改根节点
}
void rrotate(int &k){ // 右旋
    int t=l[k]; // 这将是新根节点
    l[k]=r[t]; // 原根节点的左节点是新根节点的右节点
    r[t]=k; //  新根节点的右节点是原根节点
    sz[t]=sz[k]; // 更新新根节点的 sz
    maintain(k); // 更新原根节点的 sz
    k = t; // 引用，修改根节点
}

我们注意到两个事实：1. 左旋和右旋代码是完全一样的，只不过 l 和 r 完全反过来了。 2. 我们对于根节点采用的是引用传值，这么做的目的是在每次进行修改操作时，都能够直接修改旋转的节点。

插入

k 代表目前插入的节点。如果这个节点不存在，那么就创建一个新的节点，然后 rnk 赋一个随机值。

不然先将子树大小增加 1，然后将 v 和当前节点的权值进行比较：

  • 如果 v 与当前节点相等，直接增加 cnt。
  • 如果 v 小于当前节点的权值，往左子树插入。
  • 如果 v 大于当前节点的权值，往右子树插入。

不要忘记我们要维护堆的性质，所以说每当修改后的节点的子树的 rnk 比当前节点小，我们就要把它旋转上来。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
void insert(int &k,int v){
    if(!k){ // 目前节点不存在
        k=++m;
        val[k] = v; pri[k] = rand();
        cnt[k] = 1; sz[k] = 1;
        return;
    }
    sz[k]++; // 先将子树的大小增加 1
    if(v==val[k]){ // 1. v = val[k] 原地修改
        cnt[k]++;
    } else if(v<val[k]){ // 2. v<val[k] 插入左子树
        insert(l[k],v);
        // 左子树 rnk 小了，右旋把左子节点转上来
        if(pri[l[k]]<pri[k])rrotate(k); 
    } else{ // 3. v>val[k] 插入右子树
        insert(r[k],v);
        // 右子树 rnk 小了，左旋把左子节点转上来
        if(pri[r[k]]<pri[k])lrotate(k); 
    }
}

删除

还是分三种情况讨论：

  • 如果 v 小于当前节点的权值，往左子树找，将当前节点的 sz 减 1。
  • 如果 v 大于当前节点的权值，往右子树找，将当前节点的 sz 减 1。
  • 如果 v 与当前节点相等：
    1. cnt 大于 1，直接将当前节点的 sz 减 1，然后 cnt--，结束。
    2. cnt 等于 1，这下就麻烦了，我们考虑谁来代替这个节点。
      • 最多有一个子节点：
        直接把那个子节点拿来代替即可，若没有子节点，直接把这个节点删了就可以。
      • 两个子节点都有：
        我们把 rnk 更小的那个子节点旋转上去（这样仍然保持堆的性质），然后递归删除，直到我们目前被删除的这个节点被旋转到最底层才会被真正删除。

注意维护好 sz 数组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
void del(int &k, int v){
    if(v < val[k]){
        del(l[k],v);
        sz[k]--;
    }else if(v>val[k]){
        del(r[k],v);
        sz[k]--;
    }else{
        if(cnt[k]>1){ // 不止一个，直接 cnt -1
            sz[k]--;
            cnt[k]--;
            return;
        }
        // 只有一个的情况

        if(!l[k] || !r[k]){ // 最多一个子树
            k = l[k]+r[k]; // 直接把下面那个节点弄上来
            return;
        }
        // 两个子树都有
        if(pri[l[k]]<pri[r[k]]){ // 左边小，右旋转上来
            rrotate(k);
            del(k,v); // 递归删除，子树的大小会在递归然后找子树的过程中 -1
        }else{
            lrotate(k);
            del(k,v);
        }
    }
}

查询排名

修改操作已经结束了，所以后面不会有旋转了！后面都是非常简单的搜索二叉树操作。

一样，我们还是分三类讨论：

  • 如果 v 与当前节点相等：
    排名就是左子树大小加 1。
  • 如果 v 小于当前节点的权值：
    往左子树找即可。（如果左子树是空的，说明这个数是最小的，返回 1）
  • 如果 v 大于当前节点的权值：
    排名先加上左子树的大小以及当前节点的 cnt，然后往右子树找。（如果没有右子树，直接返回当前节点的子树大小加 1）

我们这个写法可以查询不存在的数的排名。

1
2
3
4
5
6
7
8
9
10
int rnk(int k,int v){
    if(val[k]==v)return sz[l[k]]+1;
    if(v<val[k]){
        if(l[k])return rnk(l[k],v);
        else return 1;
    }

    if(r[k])return sz[l[k]]+cnt[k]+rnk(r[k],v);
    else return sz[k]+1;
}

查询第 k 大

分三类：

  1. 如果 k 小于等于左子树的大小，要查询的节点在左子树。
  2. 如果 k 大于左子树的大小，并小于等于左子树的大小加当前节点的个数，要查询的节点就是当前节点。
  3. 否则，要查询的节点在右子树，注意：我们要将排名减去左子树的大小加当前节点的个数，将排名转化为在右子树内的排名。
1
2
3
4
5
int kth(int k,int x){
    if(x<=sz[l[k]])return kth(l[k],x);
    else if(x<=sz[l[k]]+cnt[k])return val[k];
    return kth(r[k],x-sz[l[k]]-cnt[k]);
}

前驱

  1. 如果当前节点的值大于等于要查询的值，往左子树找。
  2. 现在当前节点值已经比要查询的值大了，但是我们不确定是否是最大的。因此，我们先更新一下全局变量 ans，然后再往右子树找（注意，ans 可以直接赋值，而不是用 max 操作，因为后面的赋值一定更大）。
1
2
3
4
5
6
7
8
9
int ans;
void pre(int k,int x){
    if(!k)return;
    if(val[k]>=x)pre(l[k],x);
    else{
        ans=val[k];
        pre(r[k],x);
    }
}

后继

和前驱一个思路。

1
2
3
4
5
6
7
8
void sub(int k,int x){
    if(!k)return;
    if(val[k]<=x)sub(r[k],x);
    else{
        ans=val[k];
        sub(l[k],x);
    }
}

完整代码

洛谷P3369 【模板】普通平衡树

如果你想要把这个代码，复制到其他题目里面用，注意要修改数组的大小！这里数组只开了 10510^5。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
#include<bits/stdc++.h>
using namespace std;

int read(){
    int f=1,x=0;char c=getchar();
    while(!isdigit(c)){if(c=='-')f=-1;c=getchar();}
    while(isdigit(c)){x=x*10+c-'0';c=getchar();}
    return x*f;
}

const int N = 1e5+5;
int l[N],r[N]; // 左右儿子
int val[N],pri[N]; // 值和随机的优先级
int cnt[N],sz[N]; // 重复次数,子树大小
int root; // 树根
int m; // 节点个数

void maintain(int i){ // 维护 sz 的大小
    sz[i] = cnt[i] + sz[l[i]] + sz[r[i]];
}

void lrotate(int &k){ // 左旋
    int t=r[k]; // 这将是新根节点
    r[k]=l[t]; // 原根节点的右节点是新根节点的左节点
    l[t]=k; // 新根节点的左节点是原根节点
    sz[t]=sz[k]; // 更新新根节点的 sz
    maintain(k); // 更新原根节点的 sz
    k = t; // 引用，修改根节点
}
void rrotate(int &k){ // 右旋
    int t=l[k]; // 这将是新根节点
    l[k]=r[t]; // 原根节点的左节点是新根节点的右节点
    r[t]=k; //  新根节点的右节点是原根节点
    sz[t]=sz[k]; // 更新新根节点的 sz
    maintain(k); // 更新原根节点的 sz
    k = t; // 引用，修改根节点
}

void insert(int &k,int v){
    if(!k){ // 目前节点不存在
        k=++m;
        val[k] = v; pri[k] = rand();
        cnt[k] = 1; sz[k] = 1;
        return;
    }
    sz[k]++; // 先将子树的大小增加 1
    if(v==val[k]){ // 1. v = val[k] 原地修改
        cnt[k]++;
    } else if(v<val[k]){ // 2. v<val[k] 插入左子树
        insert(l[k],v);
        // 左子树 rnk 小了，右旋把左子节点转上来
        if(pri[l[k]]<pri[k])rrotate(k); 
    } else{ // 3. v>val[k] 插入右子树
        insert(r[k],v);
        // 右子树 rnk 小了，左旋把左子节点转上来
        if(pri[r[k]]<pri[k])lrotate(k); 
    }
}

void del(int &k, int v){
    if(v < val[k]){
        del(l[k],v);
        sz[k]--;
    }else if(v>val[k]){
        del(r[k],v);
        sz[k]--;
    }else{
        if(cnt[k]>1){ // 不止一个，直接 cnt -1
            sz[k]--;
            cnt[k]--;
            return;
        }
        // 只有一个的情况

        if(!l[k] || !r[k]){ // 最多一个子树
            k = l[k]+r[k]; // 直接把下面那个节点弄上来
            return;
        }
        // 两个子树都有
        if(pri[l[k]]<pri[r[k]]){ // 左边小，右旋转上来
            rrotate(k);
            del(k,v); // 递归删除，子树的大小会在递归然后找子树的过程中 -1
        }else{
            lrotate(k);
            del(k,v);
        }
    }
}

int rnk(int k,int v){
    if(val[k]==v)return sz[l[k]]+1;
    if(v<val[k]){
        if(l[k])return rnk(l[k],v);
        else return 1;
    }

    if(r[k])return sz[l[k]]+cnt[k]+rnk(r[k],v);
    else return sz[k]+1;
}

int kth(int k,int x){
    if(x<=sz[l[k]])return kth(l[k],x);
    else if(x<=sz[l[k]]+cnt[k])return val[k];
    return kth(r[k],x-sz[l[k]]-cnt[k]);
}

int ans;
void pre(int k,int x){
    if(!k)return;
    if(val[k]>=x)pre(l[k],x);
    else{
        ans=val[k];
        pre(r[k],x);
    }
}
void sub(int k,int x){
    if(!k)return;
    if(val[k]<=x)sub(r[k],x);
    else{
        ans=val[k];
        sub(l[k],x);
    }
}

int main(){
    int n;
    srand(114514);
    scanf("%d", &n);
    int opt, x;
    for (int i = 1; i <= n; i++){
        scanf("%d%d", &opt, &x);
        if (opt == 1) insert(root, x);
        else if (opt == 2) del(root, x);
        else if (opt == 3) printf("%d\n", rnk(root, x));
        else if (opt == 4) printf("%d\n", kth(root, x));
        else if (opt == 5){
            pre(root, x);
            printf("%d\n", ans);
        }
        else if (opt == 6){
            sub(root, x);
            printf("%d\n", ans);
        }
    }
    return 0;
}