学习笔记：浅谈 pb_ds 库中的平衡树

· · 算法·理论

写在前面：作者一稿时曾随机说话，并且一些表述不够严谨，非常感谢@Seauy 老师不厌其烦地认真审核和严谨求证，辛苦了！

前言

pb_ds 库全称 Policy-Based Data Structures（基于策略的数据结构），其中一共封装了平衡树、字典树、哈希表和堆四种容器，本文主要介绍使用 pb_ds 库维护平衡树。

值得一提的是，在 2021 年以前 CCF 是不允许使用以下划线开头的库函数或宏的，好在 2021 年 CCF 发布了 关于NOI系列活动中编程语言使用限制的补充说明 ，其中明确指出了允许使用以下划线开头的库函数或宏。

本文的编写参考了 GCC 官方文档上有关于 pb_ds 库的介绍。

头文件 & 命名空间

pb_ds 库主要包含在：

ext/pb_ds/assoc_container.hpp
ext/pb_ds/tree_policy.hpp

这两个头文件中，当然，也可以使用

bits/extc++.h

万能头，<bits/extc++.h> 包含了 bits/stdc++.h 和 pb_ds 库，简直是我这种健忘症福音。

pb_ds 库的命名空间为 __gnu_pbds，但是如果 using namespace __gnu_pbds; 可能会导致 priority_queue 和 STL（位于 std）引起命名冲突，在实际使用时需谨慎。

定义

pb_ds 库的平衡树定义模板为

tree<Key, Mapped, Cmp_Fn, Tag, Node_Update, Allocator> name;

参数说明

模板参数 作用说明 取值
Key 存储的键/元素类型 自定义类型
Mapped 映射值类型，类似 set 填 null_type；类似 map 填 V 类型 null_type 等
Cmp_Fn 键的比较函子，定义元素在树中的排序规则 less<Key> 或 greater<Key>
Tag 底层数据结构实现，即平衡树的类型 rb_tree_tag（红黑树）、splay_tree_tag（伸展树）、ov_tree_tag（有序向量树）
Node_Update 节点更新策略 tree_order_statistics_node_update（启用排名功能）、null_type（默认不启用）
Allocator 内存空间分配器 allocator<char>（默认，通常省略）

例子

在 OI 竞赛中最常用的定义模板是

tree<PII , null_type , less<PII> , rb_tree_tag , tree_order_statistics_node_update> a;

（#define PII pair<int,int>）

使用 PII 作为元素类型的原因主要是因为其类似于 STL 中的 set 自动去重，留一维记录 id 即可避免去重。

至于使用 rb_tree_tag 作为参数的原因是因为红黑树常数小，相较其他平衡树有着常数上的优势。

tree_order_statistics_node_update 则是为了让平衡树支持排名。

成员函数

函数 作用 返回值 时间复杂度
size() 返回当前树中元素个数 size_type O(1)
empty() 判断树是否为空 bool ^
max_size() 返回理论最大容量（极少使用） size_type ^
clear() 清空树中所有元素 void ^
insert(value) 插入元素（或键值对） pair<iterator, bool> O(\log n)
insert(first, last) 插入迭代器区间内的所有元素 void ^
erase(iterator) 删除迭代器指向的元素 ^ ^
erase(first, last) 删除迭代器区间 [\text{first}, \text{last}) 内的元素 ^ ^
erase(key) 删除所有等于 \text{key} 的元素 size_type（删除个数，set 中为 0 或 1） ^
find(key) 查找键为 \text{key} 的元素 迭代器（找不到返回 end()） ^
lower_bound(key) 返回第一个 \ge \text{key} 的元素的迭代器 迭代器 ^
upper_bound(key) 返回第一个 > \text{key} 的元素的迭代器 ^ ^
equal_range(key) 返回 pair<lower_bound(key), upper_bound(key)> pair<iterator, iterator> ^
order_of_key(key) 返回树中严格小于 \text{key} 的元素个数 size_type（相当于排名 - 1） ^
find_by_order(index) 返回排名为 \text{index}（从 0 开始）的元素的迭代器 迭代器（若 \text{index} \ge \text{size}，返回 end()） ^

注：使用 order_of_key(key) 和 find_by_order(index) 需要启用 tree_order_statistics_node_update。

一些不常用的成员函数

函数 作用 复杂度 注
join(other) 将 other 树中的所有元素合并到当前树中，合并后 other 被清空 理论 O(\log n) 但性能不佳[^1] 两棵树的值域不能有交集（即所有元素必须严格小于或大于另一棵树的所有元素）
split(key, other) 将当前树中所有 > \text{key} 的元素移动到 other 树中，\le \text{key} 的元素留在当前树 理论 O(\log^2 n)，但实测可能退化为 O(n)[^2] 要求 other 在调用前必须为空

注：join 不检查值域是否相交，若相交则为 UB（未定义行为），请务必确保两棵树的值域完全不相交。

一些迭代器

函数 作用
begin() 返回正向迭代器起始（指向最小值）
end() 返回正向迭代器末尾
rbegin() 返回反向迭代器起始（指向最大值）
rend() 返回反向迭代器末尾

其他函数

如果模板参数 Mapped 不是 null_type（即作为 map 使用），额外支持：

函数 作用
operator[](key) 访问或插入键值对（若 \text{key} 不存在则默认构造值插入）
at(key) 安全访问（若 \text{key} 不存在则抛出 out_of_range 异常）

例题

P3871 [TJOI2010] 中位数

额啊这个原来是对顶堆吗。 维护动态中位数，支持插入和按排名查找即可。

#include <bits/extc++.h>
#define PII pair<int,int>
using namespace std;
using namespace __gnu_pbds;
typedef long long LL;
//十年OI一场空，不开long long见祖宗
tree<PII , null_type , less<PII> , rb_tree_tag , tree_order_statistics_node_update> a;
int n , m;
char s[10];
int main(){
    //freopen(".in" , "r" , stdin);
    //freopen(".out" , "w" , stdout);
    scanf("%d" , &n);
    for(int i = 1 ; i <= n ; i++){
        int x;
        scanf("%d" , &x);
        a.insert({x , i});
    }
    scanf("%d" , &m);
    getchar();
    while(m--){
        scanf("%s" , s + 1);
        if(s[1] == 'a'){
            int x;
            scanf("%d" , &x);
            a.insert({x , ++n});
        }
        else{
            auto it = a.find_by_order(n % 2 == 0 ? n / 2 - 1 : n / 2);
            printf("%d\n" , it->first);
        }
    }
    return 0;
}

P3369 【模板】普通平衡树

平衡树模板，直接看代码。

#include <bits/extc++.h>
#define PII pair<int,int>
using namespace std;
using namespace __gnu_pbds;
typedef long long LL;
//十年OI一场空，不开long long见祖宗
int n , cnt;
typedef tree<PII , null_type , less<PII> , rb_tree_tag , tree_order_statistics_node_update> Balanced_Tree;
Balanced_Tree a;
int main(){
    //freopen(".in" , "r" , stdin);
    //freopen(".out" , "w" , stdout);
    scanf("%d" , &n);
    while(n--){
        int opt , x;
        scanf("%d%d" , &opt , &x);
        if(opt == 1)
            a.insert({x , ++cnt});
        else if(opt == 2){
            auto it = a.lower_bound({x , 0});
            if(it != a.end())
                a.erase(it);
        }
        else if(opt == 3)
            printf("%d\n" , a.order_of_key({x , 0}) + 1);
        else if(opt == 4){
            auto it = a.find_by_order(x - 1);
            printf("%d\n" , it->first);
        }
        else if(opt == 5){
            auto it = a.lower_bound({x , 0});
            it--;
            printf("%d\n" , it->first);
        }
        else{
            auto it = a.lower_bound({x + 1 , 0});
            printf("%d\n" , it->first);
        }
    }
    return 0;
}

P6136 【模板】普通平衡树（数据加强版）

和 P3369 【模板】普通平衡树 没啥区别，证明一下时间复杂度。

#include <bits/extc++.h>
#define PII pair<int,int>
using namespace std;
using namespace __gnu_pbds;
typedef long long LL;
//十年OI一场空，不开long long见祖宗
int n , m , cnt , last , ans;
typedef tree<PII , null_type , less<PII> , rb_tree_tag , tree_order_statistics_node_update> Balanced_Tree;
Balanced_Tree a;
int main(){
    //freopen(".in" , "r" , stdin);
    //freopen(".out" , "w" , stdout);
    scanf("%d%d" , &n , &m);
    for(int i = 1 ; i <= n ; ++i){
        int x;
        scanf("%d" , &x);
        a.insert({x , ++cnt});
    }
    while(m--){
        int opt , x;
        scanf("%d%d" , &opt , &x);
        x ^= last;
        if(opt == 1)
            a.insert({x , ++cnt});
        else if(opt == 2){
            auto it = a.lower_bound({x , 0});
            a.erase(it);
        }
        else if(opt == 3)
            last = a.order_of_key({x , 0}) + 1;
        else if(opt == 4){
            auto it = a.find_by_order(x - 1);
            last = it->first;
        }
        else if(opt == 5){
            auto it = a.lower_bound({x , 0});
            it--;
            last = it->first;
        }
        else{
            auto it = a.lower_bound({x + 1 , 0});
            last = it->first;
        }
        if(opt >= 3)
            ans ^= last;
    }
    printf("%d" , ans);
    return 0;
}

[^1]:这里官方文档中的描述是 “These methods are efficient - red-black trees are split and joined in poly-logarithmic complexity”（翻译：这些方法是有效的，红黑树以多对数复杂度分裂和合并）。

[^2]:这里官方文档中的描述同样是多对数复杂度的，这里引用@Seauy 老师对复杂度的分析：“好像是只能分析出 O(\log^2 n) 的，split 每走到一个结点需要做一遍 split_at_node，这个函数需要从 other 的根开始遍历将当前结点与他的右子树插入到合适的位置”（再次感谢@Seauy 老师的严谨求证）。同时，曾有报告指出其底层实现可能退化到 O(n)，数据量大时请谨慎使用。