【模板】普通平衡树
同一道顺序统计题给出三种解法:离线坐标压缩 + Fenwick、Python 版 FHQ-Treap,以及 C++ 版 FHQ-Treap。
启发记录: 离线坐标压缩 + 树状数组即可完整实现普通平衡树的六个操作,说明顺序统计的本质是「前缀和 + 前缀选择」,不必手写平衡树。
OJ: luogu
题目 ID: P3369
难度:提高+/省选-
标签:树状数组坐标压缩有序多重集Treappythoncpp
日期: 2026-07-16 19:57
目录
形式化题目
维护一个可重集合
所有操作只依赖
解法总览
三种解法都能独立完成全部六个操作,差别只在「用什么结构维护顺序」:
- 解法一(
main.py,正式主解):输入已经把全部操作给全了,于是先离线收集所有作为数值出现的,去重排序做坐标压缩,再用 Fenwick 树维护每个值的出现次数。查询全部化成「前缀和」与「前缀选择」,是 Python 里常数最小、最不容易写错的写法。 - 解法二(
main-fhq-treap.py):真正意义上的在线平衡树。用 FHQ-Treap 的按值分裂与合并维护有序多重集,不需要提前知道任何数值,适合强制在线的场景。 - 解法三(
main.cpp):同一套 FHQ-Treap 思路的 C++ 实现。区别在于用数组内存池代替节点对象、用mt19937代替 Python 的随机数,常数更小。
两个 FHQ-Treap 解法的接口完全一致,可以先看 Python 版理解 split / merge,再看 C++ 版对照实现细节。
解法一:离线坐标压缩 + Fenwick
思路
题目虽然叫平衡树,但全部操作已经在输入中,可以先离线收集所有作为“数值”的 x 并坐标压缩。树状数组在压缩坐标上维护每个值的出现次数。
- 排名:
小于 x 的数量 + 1; - 第
k小:在 Fenwick 树上二进制提升,找最小前缀和达到k的坐标; - 前驱:小于
x的数量所对应的第count小; - 后继:小于等于
x的数量再加一所对应的元素。
重复值只改变同一坐标的计数,语义与可重集合一致。
Python 知识
sorted({value for ...})一行完成离线去重排序。bisect_left统计严格小于,bisect_right统计小于等于。- Fenwick 的
kth使用bit_length()取得最高二进制步长。 - 操作保存成元组列表,第二遍执行时无需重新解析。
代码
import sys
from bisect import bisect_left, bisect_right
# ===== 快速读入 =====
# sys.stdin.buffer.read() 一次性读入全部字节(比逐行 input() 快得多)。
# .split() 按空白切成 bytes 列表,交给 map(int, ...) 惰性转换。
# 注意 map 返回的是迭代器,只能顺着取一次,所以用 iter() 包一层后统一用 next() 取值。
data = iter(map(int, sys.stdin.buffer.read().split()))
# 第一项是操作次数 m,之后每条操作占两个数:op 和 x。
# 列表推导式一次把 m 条操作取成 (op, x) 元组;因为输入保证成对出现,这里不会越界。
operations = [(next(data), next(data)) for _ in range(next(data))]
# ===== 坐标压缩(离线)=====
# 题目叫“平衡树”,但所有操作已经在输入里给全了,可以先纵览全局:
# 收集所有出现过的数值 x(操作 4 的 x 是排名,不是数值,要排除),
# 用集合去重后排序,得到“值 -> 下标”的映射。之后树状数组只需在这些下标上工作,
# 长度从 1e9 级别的值域压缩到 O(m)。
coordinates = sorted({value for operation, value in operations if operation != 4})
# 树状数组本体,1-indexed,多开一格让 sizes+1 个元素都进得来。
tree = [0] * (len(coordinates) + 1)
def add(index, delta):
"""把坐标 index(0-based,调用方传入)的出现次数增加 delta。"""
index += 1 # Fenwick 内部从 1 开始,外部统一用 0-based 坐标
while index < len(tree):
tree[index] += delta
index += index & -index # 跳到下一个被该点覆盖的位置(加最低位的 1)
def prefix(end):
"""返回区间 [0, end) 内的总计数,即前 end 个坐标的出现次数之和。"""
result = 0
while end:
result += tree[end]
end -= end & -end # 抹掉最低位的 1,跳到上一层区间
return result
def kth(rank):
"""返回整棵树上第 rank 小(1-based)的数值。"""
index = 0 # 已经走过的前缀长度(也是 1-based 坐标)
# bit_length() 取二进制位数,向下取到最大的 2 的幂作为第一步长,
# 这正是树状数组能直接二分(二进制提升)的原因。
step = 1 << (len(tree).bit_length() - 1)
while step:
target = index + step
# 若 target 这一整段的前缀和仍小于 rank,说明答案不在这段里,整段跳过
# 为什么这里用到了tree[target]? 这里是BITree的二进制跳跃的最大的性质: 正好是跳跃的区间
if target < len(tree) and tree[target] < rank:
index = target
rank -= tree[target]
step >>= 1
# 循环结束时 index 是最后一个前缀和 < rank 的坐标,加 1 才是答案坐标;
# 函数开头 index = 0,所以 coordinates[index] 恰好取到 0-based 的答案。
return coordinates[index]
# ===== 第二遍处理操作 =====
# 结果先存进列表,最后用 "\n".join 一次输出:
# 比在循环里反复 print() 少很多次系统调用,是 Python 提速的常用技巧。
answers = []
for operation, value in operations:
if operation == 1: # 插入 x:该值计数 +1
add(bisect_left(coordinates, value), 1)
elif operation == 2: # 删除一个 x:该值计数 -1
add(bisect_left(coordinates, value), -1)
elif operation == 3: # 查询 x 的排名 = 小于 x 的个数 + 1
# bisect_left 返回第一个 >= value 的下标,正好是“严格小于 value”的个数
answers.append(str(prefix(bisect_left(coordinates, value)) + 1))
elif operation == 4: # 查询第 x 小
answers.append(str(kth(value)))
elif operation == 5: # 严格前驱:小于 x 的个数,对应第 count 小
answers.append(str(kth(prefix(bisect_left(coordinates, value)))))
else: # operation == 6,严格后继:小于等于 x 的个数 + 1,对应那个排名
# bisect_right 返回第一个 > value 的下标,即“小于等于 value”的个数
answers.append(str(kth(prefix(bisect_right(coordinates, value)) + 1)))
print("\n".join(answers))复杂度
设操作数为
解法二:FHQ-Treap(Python)
思路
FHQ-Treap 只有两个核心操作,其余六个接口都由它们拼出来:
split(u, val)把树按值切成 (<= val)与(> val)两棵,递归下去后顺着原路径改接儿子并更新子树大小;merge(x, y)要求中所有值不大于 中所有值,比较两棵树根的随机优先级,优先级大的当父节点,继续往下递归合并。
有了这两个操作,六种接口都很短:
- 插入:按
val切成两半,把新节点夹在中间合并回去; - 删除一个:按
val切成(<= val),再按val - 1把左边切成(< val)和(== val),丢掉后者的根(它的左右子树合并起来),再把三段拼回; - 排名 / 前驱 / 后继:不改变树结构,直接按 BST 性质走一遍,靠「左子树大小 + 1」累计比
val小的元素个数,与前缀和思路完全等价; - 第
小:同样从根往下走,比较 与左子树大小决定去哪边。
随机优先级决定树高。 优先级是每个节点独立随机生成的,merge 只在两者之间比大小,所以树的形状等价于对节点随机建堆,期望高度 random.getrandbits(64) 做种子而不是写死常量,避免出题人针对固定种子构造让树退化的数据。
Python 知识
Node用__slots__声明val/pri/size/l/r,省掉__dict__,节点数量到时内存和访问速度都有明显收益。 split/merge是递归的,需要sys.setrecursionlimit留出余量。- 用
data = iter(map(int, sys.stdin.buffer.read().split()))一次性读入后顺序next()取值,避免逐行input()。 random.Random(seed)与全局random隔离,方便需要时复现随机序列。
代码
import random
import sys
from typing import Optional, Tuple, Any
# FHQ-Treap 的 split / merge 都是递归实现的。
# 随机优先级保证树高期望为 O(log n),但为了保险仍把递归上限调高。
sys.setrecursionlimit(300000)
class Node:
"""Treap 节点"""
__slots__ = ("val", "pri", "size", "l", "r")
def __init__(self, val: Any, pri: int):
self.val = val
self.pri = pri
self.size = 1
self.l: Optional["Node"] = None
self.r: Optional["Node"] = None
class FHQTreap:
"""
FHQ-Treap(无旋平衡树)
基于按值分裂 (split) 与合并 (merge) 维护动态有序集合。
支持:插入、删除单次出现、按值查排名、查第 k 小、查前驱、查后继。
所有基本操作的期望时间复杂度均为 O(log n)。
"""
def __init__(self, seed: Optional[int] = 233):
self.root: Optional[Node] = None
self.rng = random.Random(seed)
def _size(self, u: Optional[Node]) -> int:
return u.size if u is not None else 0
def _push_up(self, u: Node) -> None:
"""由左右子树大小更新当前节点子树大小"""
u.size = self._size(u.l) + self._size(u.r) + 1
def _new_node(self, val: Any) -> Node:
return Node(val, self.rng.randint(1, 2**31 - 1))
def split(self, u: Optional[Node], val: Any) -> Tuple[Optional[Node], Optional[Node]]:
"""
按数值 val 分裂:
- 左树 x:包含所有节点值 <= val 的节点
- 右树 y:包含所有节点值 > val 的节点
"""
if u is None:
return None, None
if u.val <= val:
x, y = self.split(u.r, val)
u.r = x
self._push_up(u)
return u, y
else:
x, y = self.split(u.l, val)
u.l = y
self._push_up(u)
return x, u
def merge(self, x: Optional[Node], y: Optional[Node]) -> Optional[Node]:
"""
合并两棵树 x 和 y:
前提:x 中所有节点的值 <= y 中所有节点的值
依据节点的随机优先级保持大根堆性质
"""
if x is None or y is None:
return x if y is None else y
if x.pri > y.pri:
x.r = self.merge(x.r, y)
self._push_up(x)
return x
else:
y.l = self.merge(x, y.l)
self._push_up(y)
return y
def insert(self, val: Any) -> None:
"""插入一个数值 val"""
x, y = self.split(self.root, val)
node = self._new_node(val)
self.root = self.merge(self.merge(x, node), y)
def delete(self, val: Any) -> None:
"""
删除一个数值等于 val 的节点(若存在多个同值节点仅删除其中一个)。
通过将集合切为 (< val)、(== val)、(> val) 三部分,丢弃 (== val) 的一个节点后拼回。
"""
x, z = self.split(self.root, val)
x, y = self.split(x, val - 1)
if y is not None:
# 丢弃 y 的根节点,将其左右子树合并
y = self.merge(y.l, y.r)
self.root = self.merge(self.merge(x, y), z)
def rank(self, val: Any) -> int:
"""
查询数值 val 在集合中的排名(小于 val 的元素个数 + 1)。
采用 BST 遍历,常数小于 split/merge 且无需改变树结构。
"""
u = self.root
ans = 0
while u is not None:
if u.val < val:
ans += self._size(u.l) + 1
u = u.r
else:
u = u.l
return ans + 1
def kth(self, k: int) -> Any:
"""
查询集合中第 k 小的元素(1-based)。
若 k 超出合法范围 [1, size()] 则抛出 IndexError。
"""
if not (1 <= k <= self.size()):
raise IndexError(f"kth index {k} out of range (size={self.size()})")
u = self.root
while u is not None:
l_sz = self._size(u.l)
if k <= l_sz:
u = u.l
elif k == l_sz + 1:
return u.val
else:
k -= l_sz + 1
u = u.r
raise IndexError("kth not found")
def pre(self, val: Any) -> Optional[Any]:
"""
查询 val 的前驱(小于 val 的最大值)。
若不存在严格小于 val 的值,返回 None。
"""
u = self.root
ans = None
while u is not None:
if u.val < val:
ans = u.val
u = u.r
else:
u = u.l
return ans
def succ(self, val: Any) -> Optional[Any]:
"""
查询 val 的后继(大于 val 的最小值)。
若不存在严格大于 val 的值,返回 None。
"""
u = self.root
ans = None
while u is not None:
if u.val > val:
ans = u.val
u = u.l
else:
u = u.r
return ans
def size(self) -> int:
"""返回集合中的节点总数"""
return self._size(self.root)
def __len__(self) -> int:
return self.size()
def empty(self) -> bool:
"""判断集合是否为空"""
return self.root is None
def clear(self) -> None:
"""清空平衡树"""
self.root = None
def main() -> None:
# 一次性读入全部输入并按空白切分,再惰性转成整数。
# 用迭代器顺序取值,比循环里反复 input() / split() 快很多。
data = iter(map(int, sys.stdin.buffer.read().split()))
n = next(data)
# 用系统熵做随机种子:如果固定种子,出题人可能针对性地构造让树退化的数据。
tree = FHQTreap(random.getrandbits(64))
answers = []
for _ in range(n):
opt = next(data)
x = next(data)
if opt == 1: # 插入 x
tree.insert(x)
elif opt == 2: # 删除一个 x
tree.delete(x)
elif opt == 3: # x 的排名 = 小于 x 的个数 + 1
answers.append(str(tree.rank(x)))
elif opt == 4: # 第 x 小
answers.append(str(tree.kth(x)))
elif opt == 5: # 前驱:小于 x 的最大值
answers.append(str(tree.pre(x)))
else: # opt == 6,后继:大于 x 的最小值
answers.append(str(tree.succ(x)))
# 结果统一收集后一次输出,避免大量 print() 的系统调用开销。
print("\n".join(answers))
if __name__ == "__main__":
main()
复杂度
六种操作都是「一次分裂 + 若干次合并」或「一次 BST 下降」,期望时间均为
解法三:FHQ-Treap(C++)
思路
思路与解法二完全相同:按值分裂、按随机优先级合并,rank / kth / pre / succ 用 BST 下降实现。差异只在工程细节:
- 节点不单独
new,而是预先开一个Node tr[N]数组,用下标代替指针,tr[0]作为哨兵(左右儿子与大小全为 0),这样update不用判空; - 随机优先级用
mt19937; - 前驱、后继在没找到时分别返回
numeric_limits<T>::min()与::max(),模板参数T改成long long也能用。
代码
/**
* Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
* rbook: -> https://rbook.roj.ac.cn https://rbook2.roj.ac.cn
* rainboy的学习导航网站: https://idx.roj.ac.cn
* create_at: 2026-07-16 19:57
* update_at: 2026-07-16 19:57
*/
#include <bits/stdc++.h>
#include <random>
#include <limits>
using namespace std;
template<typename T = long long,int N = 500005>
struct FHQ
{
// ============ 数据
int root;
struct Node {
int l,r;
int size,fix;
T val;
};
// 使用内存池 替代
Node tr[N];
int tr_idx = 0;
int get() { return ++tr_idx;}
std::mt19937 rnd;
// ============ 数据
// 定义极值,用于表示“没找到”
// 如果是 long long 题目,请确保 T 是 long long
const T INF_MAX = std::numeric_limits<T>::max();
const T INF_MIN = std::numeric_limits<T>::min();
// ============ 构造函数
FHQ() {
// root = 0;
rnd.seed(233);
// 也可以用下面的方式初始化随机数种子
// rnd.seed(std::random_device()());
init();
}
// --- 1. 多组数据必备 (Clear & Init) ---
void init() {
root = 0;
tr_idx = 0;
// 哨兵清理:防止 tr[0] 脏数据污染 size 计算
tr[0].l = tr[0].r = tr[0].size = 0;
tr[0].val = 0;
}
void clear() {
init();
}
int size() const {
return tr[root].size;
}
bool empty() const {
return size() == 0;
}
// ============ 核心操作
int new_node(T v) {
int id = get();
// tr[idx].ch[0] = tr[idx].ch[1] = 0;
tr[id].l = tr[id].r = 0; // 与上面等价
tr[id].size = 1;
tr[id].fix = rnd();
tr[id].val = v;
return id;
}
void update(int u ) {
int l = tr[u].l;
int r = tr[u].r;
tr[u].size = tr[l].size + tr[r].size + 1;
}
// --- 这里的两个函数是 FHQ-Treap 的灵魂 ---
/**
* 分裂 (Split) - 按数值 v 分裂
* 含义:将树 u 切成两棵树 x 和 y。
* 功能:
* 1. 使得 x 指向 u 中所有节点值 <= v 的节点组成的树
* 2. 使得 y 指向 u 中所有节点值 > v 的节点组成的树
* x 树:所有节点值 <= v
* y 树:所有节点值 > v
* 注意:x 和 y 是通过引用传回的
* 注: 经过我的各种写法,发现最好的写法就是下面的这个,不要在修改了 by rainboy 2025-11-17
*/
void split(int u,int v,int &x,int &y) {
// 边界: 如果 u 是空树,那么 x 和 y 都是空树
if( !u ) { x = y = 0; return; }
if( tr[u].val <= v ) {
x = u;
// 剩下的问题是:u 的右子树里,可能还有 <= v 的,也可能有一部分 > v 的
// 所以我们递归处理 u 的右子树
split(tr[u].r,v,tr[u].r,y);
} else {
// 当前节点 > v,说明当前节点及其右子树都属于 y
y = u;
// 递归处理 u 的左子树
split(tr[u].l,v,x,tr[u].l);
}
update(u); // 更新 u 的 size , 因为 u 的左右子树可能发生了变化
}
/**
* 合并 (Merge)
* 含义:将两棵树 x 和 y 合并成一棵树,并返回新根。
* 前提:x 中所有节点的值 必须 <= y 中所有节点的值 (即有序)
* 技巧:利用随机优先级 pri 来决定谁做父节点,从而保持平衡
*/
int merge(int x,int y) {
// 边界其中一个是空树,空树是单位元:如果有一棵树是空的,那么合并结果就是另一棵树
if( !x || !y ) return x + y;
// 谁的优先级高,谁就做父节点(取决于你是大根堆还是小根堆,这里用大根堆)
if( tr[x].fix > tr[y].fix ) {
tr[x].r = merge(tr[x].r,y);
update(x); // 更新 x 的 size , 因为 x 的右子树可能发生了变化
return x;
} else {
tr[y].l = merge(x,tr[y].l);
update(y);
return y;
}
}
// --- 常用操作封装 (心智负担最低的写法) ---
/**
* 插入 (Insert)
* 含义:将值 v 插入到树中
* 注: 这个是最好的插入, 如果使用BST的插入,需要旋转,那就是普通Treap
*/
void insert(T v) {
int x,y;
// 先按 v 分裂,把树切成 (<=v) 和 (>v) 两部分
// 注: split <=v ,因为 新的节点放到相同的v的后面
split(root,v,x,y);
root = merge(merge(x,new_node(v)),y);
}
/**
* 删除 (Delete)
* 含义:将值 v 从树中删除
* 注: 这个是最好的删除, 与插入一样
*/
void del(T v) {
int x, y, z;
// 1. 把树按 v 分裂 -> x(<=v), z(>v)
split(root, v, x, z);
// 2. 把 x 按 v-1 分裂 -> x(<v), y(==v)
split(x, v - 1, x, y);
// 此时 y 树里全是值为 v 的节点。
// 如果要删除所有 v,直接忽略 y 即可。
// 如果只删除一个 v:
if (y) // y不是空树
{
// 合并 y 的左右孩子,相当于把 y 的根节点丢弃了
y = merge(tr[y].l, tr[y].r);
}
// 3. 全部装回去
root = merge(merge(x, y), z);
}
/**
* 查询排名 (Rank), split 风格
* 含义:查询值 v 在树中的排名
* 注意:v-1 的依赖 T 的类型,如果是浮点数,可能会有精度问题
*/
int rank_with_split(T v) {
int x,y;
split(root,v-1,x,y);
int ans = tr[x].size + 1;
root = merge(x,y);
return ans;
}
/**
* 查询排名 (Rank), bst风格,比split风格快
* 含义:查询值 v 在树中的排名
* 注意:v-1 的依赖 T 的类型,如果是浮点数,可能会有精度问题
*/
T rank(int v) const {
// 传统的 BST 查询方式, 不依赖 split/merge
int u = root;
T ans = 0; // 记录比 v 小的节点数量
while(u) {
if( tr[u].val < v) {
// 当前节点比 v 小, 左子树所有节点也都比 v 小
ans += tr[tr[u].l].size + 1;
u = tr[u].r; // 去右子树继续找
}
else { // tr[u].val >= v
// 当前节点不小于 v, 答案在左子树
u = tr[u].l;
}
}
return ans + 1; // 排名是 "比我小的个数 + 1"
}
/**
* 查询第 k 小 (Kth),bst 风格
* 含义:查询树中第 k 小的值
*/
T kth(int k) const {
int u = root;
while( true ) {
int l_size = tr[tr[u].l].size;
if( k <= l_size ) u = tr[u].l;
else if( k == l_size + 1 ) return tr[u].val;
else k -= l_size + 1,u = tr[u].r;
}
}
/**
* 查询前驱 (Predecessor)
* 含义:查询值 v 在树中的前驱 (小于v的最大值)
*/
T pre(T v) {
// T ans = -2147483647; // 假设为 int 类型的最小值, 表示未找到
T ans = std::numeric_limits<T>::min(); // 更通用的写法
int u = root;
while (u) {
if (tr[u].val < v) {
// 当前节点是 v 的一个潜在前驱, 记录下来
ans = tr[u].val;
// 尝试在右子树中寻找更大的前驱
u = tr[u].r;
} else {
// 当前节点值 >= v, 前驱一定在左子树
u = tr[u].l;
}
}
return ans;
}
/**
* 查询后继 (Successor)
* 含义:查询值 v 在树中的后继 (大于v的最小值)
*/
T succ(T v) {
int u = root;
// T ans = 2147483647; // 假设为 int 类型的最大值, 表示未找到
T ans = std::numeric_limits<T>::max(); // 更通用的写法
while (u) {
if (tr[u].val > v) {
// 当前节点是 v 的一个潜在后继, 记录下来
ans = tr[u].val;
// 尝试在左子树中寻找更小的后继
u = tr[u].l;
} else {
// 当前节点值 <= v, 后继一定在右子树
u = tr[u].r;
}
}
return ans;
}
// --- 4. STL 风格查询接口 ---
/**
* lower_bound: 寻找第一个 >= v 的值
* 返回值: 找到的值,如果没找到返回 INF_MAX
*/
T lower_bound(T v) {
return succ(v-1); // 等价于 succ(v-1)
}
/**
* upper_bound: 寻找第一个 > v 的值 (等同于 succ)
* 返回值: 找到的值,如果没找到返回 INF_MAX
*/
T upper_bound(T v) {
return succ(v);
}
};
FHQ fhq;
int main() {
int n;
std::cin >> n;
int opt,val,ans;
while (n--) {
std::cin >> opt >> val;
if( opt == 1) {
fhq.insert(val);
}
else if( opt == 2) {
fhq.del(val);
}
else if( opt == 3) {
auto ans = fhq.rank(val);
std::cout << ans <<endl;
}
else if( opt == 4) {
auto ans = fhq.kth(val);
std::cout << ans <<endl;
}
else if( opt == 5) {
auto ans = fhq.pre(val);
std::cout << ans <<endl;
}
else if( opt == 6) {
auto ans = fhq.succ(val);
std::cout << ans <<endl;
}
}
return 0;
}复杂度
与解法二相同:每种操作期望
复杂度对比
| 方案 | 维护结构 | 是否在线 | 时间复杂度 | 空间复杂度 |
|---|---|---|---|---|
main.py |
坐标压缩 + Fenwick | 否(需先读完输入) | 预排序 |
|
main-fhq-treap.py |
FHQ-Treap(节点对象) | 是 | 每次操作期望 |
|
main.cpp |
FHQ-Treap(数组内存池) | 是 | 每次操作期望 |
三者的渐进复杂度相同,真正拉开差距的是常数和适用条件:Fenwick 每步只有几行数组操作,但要求能离线读完全部输入;FHQ-Treap 允许边读边做,代价是每个节点多一次随机数开销和递归调用。
总结
OJ 输入允许离线时,不必强行在 Python 手写旋转平衡树;坐标压缩 + Fenwick 同样完整实现顺序统计接口。
而如果需要在线维护,FHQ-Treap 是比旋转 Treap 更好写的一种:把「按值分裂」和「按随机优先级合并」这两个操作想清楚,插入、删除、排名、第