Treap 平衡树

id: treap

一、 Treap 解决什么问题

维护一个动态有序多重集合，支持下面六种操作，每次期望 O(log⁡n)O(\log n)：

操作 含义
插入 把 xx 放进集合
删除 删掉一个 xx
排名 比 xx 小的数有多少个
第 kk 小 集合里排名为 kk 的数
前驱 小于 xx 的最大数
后继 大于 xx 的最小数

前置知识：[ Rbook: 二叉查找树] 和 [ Rbook: 堆]。

先看两种朴素做法为什么不行：

做法 插入 / 删除 查询 问题
有序数组 O(n)O(n) O(log⁡n)O(\log n) 插入要整体搬移
普通二叉查找树 平均 O(log⁡n)O(\log n) 平均 O(log⁡n)O(\log n) 按有序数据插入会退化成链，最坏 O(n)O(n)

普通 BST 退化的原因是：树形只由插入顺序决定，而插入顺序往往不是我们能选的——题目给的数据本来就可能有序。Treap 的做法是给每个点再加一个随机关键字，让树形由随机数决定。

二、 Treap = Tree + Heap

每个点存两个关键字：

  • val：参与 BST 比较，就是集合里的那个数；
  • fix：插入时随机生成一次，之后不再改变。

整棵树同时满足两条性质：

  1. BST 性质：左子树所有 val 小于根，右子树所有 val 大于根。相等的值不新建点，把出现次数记在 cnt 里。
  2. 小根堆性质：根的 fix 小于两个孩子的 fix（这里用最小堆，用最大堆也可以，对称）。

为什么这样就平衡？因为给定所有点的 (val, fix)，树形是唯一确定的：它等于把这些点按 fix 从小到大依次插入得到的 BST。fix 是随机的，所以这棵树和「把数据随机打乱后建 BST」是同一个分布，期望高度 O(log⁡n)O(\log n)。

一句话概括：随机 fix 把「插入顺序」这件我们控制不了的事，变成了一次随机实验。

三、 节点与信息上传

typedef struct { int val, fix, sz, cnt, c[2]; } node;
node tr[maxn];
int root = 0, idx = 0;

inline int &lc(int p) { return tr[p].c[0]; }
inline int &rc(int p) { return tr[p].c[1]; }

// 左右子树的答案都正确时，重算 p 的子树大小
inline void pushup(int p) {
    tr[p].sz = tr[lc(p)].sz + tr[rc(p)].sz + tr[p].cnt;
}
  • c[0] / c[1] 存左右孩子下标，用引用函数 lc() / rc() 包一层，后面写起来短。
  • sz 是子树里数的个数（不是点数），所以要把 cnt 加上。
  • 每个点存 cnt 而不是「相同的值建多个点」，可以让相同的数共用一个点，删除时也只需要 cnt--。

四、 旋转：唯一的结构调整

插入和删除都可能破坏堆性质，修它的手段只有一种：旋转。旋转做的事是「把孩子提上来、把根压下去」，方向只有两个。

下图是右旋（把左孩子提上来），括号里是 fix：

G 4 4(30) 2 2(10) 4->2 6 6(40) 4->6 1 1(5) 2->1 3 3(20) 2->3
查看 DOT 源码 
digraph G {
    node [shape=circle, width=0.55, fontsize=12];
    edge [arrowsize=0.6];
    4 [label="4(30)"];
    2 [label="2(10)"];
    6 [label="6(40)"];
    1 [label="1(5)"];
    3 [label="3(20)"];
    4 -> 2; 4 -> 6; 2 -> 1; 2 -> 3;
}

对以 4 为根的子树右旋后，2 成为新的根：

G 2 2(10) 4 4(30) 2->4 1 1(5) 2->1 3 3(20) 4->3 6 6(40) 4->6
查看 DOT 源码 
digraph G {
    node [shape=circle, width=0.55, fontsize=12];
    edge [arrowsize=0.6];
    2 [label="2(10)"];
    4 [label="4(30)"];
    1 [label="1(5)"];
    3 [label="3(20)"];
    6 [label="6(40)"];
    2 -> 1; 2 -> 4; 4 -> 3; 4 -> 6;
}

两条必须记住的性质：

  • 旋转不改变中序遍历。上面的图里两种形态的中序都是 1 2 3 4 6，所以 BST 性质自动保持。
  • 旋转只改变父子关系，不改变任何 val。所以旋转是「为了修堆性质而做的纯结构调整」。
// aim = 0：右旋，把左孩子提上来
// aim = 1：左旋，把右孩子提上来
void rotate(int &p, bool aim) {
    int t = tr[p].c[aim];        // 要提上来的那个孩子
    tr[p].c[aim] = tr[t].c[!aim];// 孩子的反方向子树挂回 p
    tr[t].c[!aim] = p;           // p 变成孩子的反方向孩子
    pushup(p);                   // 先更新下沉的 p
    pushup(t);                   // 再更新上升的 t
    p = t;                       // p 是引用，父节点自动指向新根
}

p 用引用传参是关键：这样 p = t 会直接改掉父节点里存的儿子下标，不需要额外记录父指针。

五、 插入

和普通 BST 一样递归下行，找到空位置建点；回溯时检查孩子的 fix 是否比当前点小，小就把孩子转上来。

void _ins(int &p, int v) {
    if (!p) {                    // 空位置：新建一个点
        p = ++idx;
        tr[p] = {v, rnd(), 1, 1, {0, 0}};
        return;
    }
    tr[p].sz++;                  // v 一定落在 p 的子树里
    if (tr[p].val == v) {        // 已经存在：只加计数
        tr[p].cnt++;
        return;
    }
    bool aim = v > tr[p].val;    // 大往右(1)，小往左(0)
    _ins(tr[p].c[aim], v);
    if (tr[tr[p].c[aim]].fix < tr[p].fix)  // 孩子 fix 更小，违反小根堆
        rotate(p, aim);                     // 把孩子转上来
}

注意 tr[p].sz++ 放在「相等」判断之前：相等的值也会让子树里的个数加一。而新建点的分支不经过 sz++，它自己在初始化时就是 sz = 1。

六、 删除

先按 BST 找到 val == v 的点，然后分三种情况：

  1. cnt > 1：还有相同的数，cnt--、sz-- 即可。
  2. 至多一个孩子：直接用唯一的孩子（或空）顶替自己。
  3. 两个孩子：不能直接删，否则树会断成两棵。做法是把 fix 较小的那个孩子转上来，自己下沉一层，然后继续在下面删自己——一直转到自己变成第 1、2 种情况为止。
void _del(int &p, int v) {
    if (!p) return;
    if (tr[p].val == v) {
        if (tr[p].cnt > 1) {          // 情况 1
            tr[p].cnt--;
            tr[p].sz--;
        }
        else if (!lc(p) || !rc(p)) {  // 情况 2
            p = lc(p) + rc(p);
        }
        else {                        // 情况 3
            // 把 fix 更小的孩子转上来：左孩子 fix 更大就提右孩子
            rotate(p, tr[lc(p)].fix > tr[rc(p)].fix);
            _del(p, v);               // p 已经是新的子树根，继续往下删
        }
    }
    else {
        tr[p].sz--;
        _del(tr[p].c[v > tr[p].val], v);
    }
}

这里的方向是全文最容易写错的一行。写成 tr[lc(p)].fix < tr[rc(p)].fix 时，六种操作的答案依然全对（BST 性质没坏），但堆性质被破坏，树会慢慢退化，大数据直接超时。验证方法见第九节。

七、 四个查询

四个操作都只需要从根往下走一条链。

7.1 排名

排名定义为「比 vv 小的数的个数 + 1」。往下走时，每次往右走，就要把「左子树全部 + 当前点的 cnt」计入答案。

int Rank(int v) {
    int p = root, ans = 1;
    while (tr[p].val != v && tr[p].c[v > tr[p].val]) {
        if (v > tr[p].val) ans += tr[lc(p)].sz + tr[p].cnt;
        p = tr[p].c[v > tr[p].val];
    }
    // 走到 v 或走到空：补上「比 v 小的左子树」
    // 如果停下的点比 v 小（v 不存在时），还要加上它的 cnt
    return ans + tr[lc(p)].sz + (tr[p].val < v ? tr[p].cnt : 0);
}

7.2 第 k 小

设左子树大小为 lsz、当前点计数为 cnt：

  • lsz >= k：答案在左子树；
  • lsz + cnt < k：答案在右子树，先把这两块减掉；
  • 否则答案就是当前点。
int at(int k) {
    int p = root;
    while (p) {
        int lsz = tr[lc(p)].sz;
        if (lsz >= k) p = lc(p);
        else if (lsz + tr[p].cnt < k) k -= lsz + tr[p].cnt, p = rc(p);
        else break;
    }
    return p;
}

7.3 前驱与后继

有了「排名」和「第 kk 小」，前驱后继不需要单独写：

  • 前驱 = 第 Rank(v) - 1 小的数；
  • 后继 = 第 Rank(v + 1) 小的数。因为 Rank(v + 1) 统计的是「小于等于 vv 的个数 + 1」，正好是第一个大于 vv 的位置。
int pre(int v)  { return tr[at(Rank(v) - 1)].val; }
int next(int v) { return tr[at(Rank(v + 1))].val; }

这个写法依赖「前驱/后继存在」：vv 是最小值时 at(0) 会返回 0 号点。模板题保证查询合法；如果题目不保证，要先判 Rank(v) > 1 和后继是否存在。

八、 完整模板

const int maxn = 1e6+5,maxe = 1e6+5; //点与边的数量

struct Random {
    random_device rd;
    mt19937 engine{rd()};
    uniform_int_distribution<> dis;
    Random(){}
    Random(int l,int r){ dis = uniform_int_distribution<> (l,r); }
    int operator()(){ return dis(engine); }
};

struct Treap {
    typedef struct { int val,fix,sz,cnt,c[2];} node;
    node tr[maxn];
    Random rnd;
    int root,idx=0,top=0,s[maxn];
    inline int &lc(int p){ return tr[p].c[0];}
    inline int &rc(int p){ return tr[p].c[1];}
    inline int &c(int p,int i){ return tr[p].c[i];}
    inline int &fix(int p){return tr[p].fix;}
    inline int &sz(int p){return tr[p].sz;}
    inline int &cn(int p){return tr[p].cnt;}
    inline int &va(int p){return tr[p].val;}

    // 建立一个新的节点
    inline void New(int &p,int v){
        tr[p = ++idx] = {v,rnd(),1,1,{0,0}};
    }
    // 向上更新
    inline void pushup(int p){
        sz(p) = sz(lc(p)) + sz(rc(p)) + cn(p);
    }

    // aim == 0 把左孩子给提上来 表示右旋 
    // aim == 1 把右孩子给提上来 表示左旋
    void rotate(int &p,bool aim){
        int t = c(p,aim);       //t 对应的孩子
        c(p,aim) = c(t,!aim);   //指向孩子的反孩子
        c(t,!aim) = p;          //孩子的反孩子指向自己
        pushup(p);pushup(t);    //更新
        p = t;                  //指向新的根
    }

    //插入
    void _ins(int &p,int v){
        if( !p ) {
            New(p,v);
            return;
        }
        sz(p)++;
        if( va(p) == v) cn(p)++;
        else {
            bool aim = v> va(p);
            _ins(c(p,aim),v);
            if( fix(c(p,aim)) < fix(p) ) //孩子的fix值比自己小
                rotate(p, aim);
        }
    }
    void ins(int v){
        _ins(root,v);
    }

    //删除
    void _del(int &p,int v){
        if( !p) return;
        if(va(p) == v){
            if( cn(p) > 1) cn(p)--,sz(p)--;
            else if(!lc(p) || !rc(p) ) p = lc(p) +rc(p); //只有一个孩子存在
            else { //两个孩子都存在
                //把 fix 更小的那个孩子提上来，自己下沉一层
                //aim=0 右旋提左孩子；aim=1 左旋提右孩子
                rotate(p, fix(lc(p)) > fix(rc(p)));
                _del(p,v);//把提上来的点 sz--
            }
        }
        else {
            sz(p)--;
            _del(c(p,v>va(p)),v);
        }
    }
    void del(int v){ _del(root,v);}

    // 排名
    // 每一次向右走，排名都会增加
    // 边界1: va(p) != p 但已经不向下走了，
    //          还要加上lc(p)的数量
    //          如果va(p) < v 还加上cn(p)
    // 边界2: va(p) == v 还要加上lc(p)的数量
    int Rank(int v){ //得到v的排名
        int p = root , ans = 1;
        while( va(p) != v && c(p,v>va(p)) ){
            if( v > va(p)) ans += sz(lc(p)) + cn(p);
            p = c(p,v > va(p));
        }
        return ans + sz(lc(p)) + ( va(p) < v ? cn(p) : 0);
    }
    // 得到排名为k的值
    // 排名树的数量是 sz(lc(p))
    // 排名为k的数在左子树上： sz(lc(p)) >= k
    // 排名为k的数在右子树上： sz(lc(p)) + cn(p) < k
    // 排名为k的数在当前点上： 不满足上面的两个条件
    int at(int k){
        int p = root;
        while( p ){
            if( sz(lc(p)) >= k ) p = lc(p);
            else if ( sz(lc(p)) + cn(p) < k)
                k -= sz(lc(p)) + cn(p), p = rc(p);
            else break; }
        return p;
    }
    // pre  前趋
    // 得到v的排名k
    // 得到排名为k-1的数
    // 1 2 2 2 5
    // 1 2     5
    // rank(5) == rank(4) == rank(3)
    // at(rank(5-1)) = 5
    inline int pre(int v){
        return va(at(Rank(v)-1));
    }
    // next 后继
    // 得到的v+1的排名，
    // 1 2 2 2 5
    // 1 2     5 rank 显然应该Rank(v+1)
    // rank(2)+1 = 5
    // rank(3)+1 = 6
    // rank(4)+1 = 6
    // Rank(2) = 2 Rank(3) = 5 at(Rank(2)+1) = 2
    inline int next(int v){
        return va(at(Rank(v+1)));
    }
} treap;

模板用 mt19937 生成 fix，六种操作封装在 Treap 结构体里，多组数据时重新构造即可。

九、 易错点

  1. 删除时旋转方向反了。这是最隐蔽的错误：答案全对，只是树不再平衡。检查办法是在本地随机造操作、和自己的 std::multiset 对拍，同时遍历整棵树验证「每个孩子的 fix 都不小于父亲的 fix」。这个断言不成立就说明方向错了。
  2. 用 rand() 生成 fix。rand() 的最大值通常只有 32767，点数一多就大量重复，堆性质形同虚设；而且 rand() 在某些平台很慢。用 mt19937，或者手写一个 xorshift。
  3. sz 与 cnt 的更新时机。插入时相等的值也要 sz++；删除时先动 cnt 再动 sz；旋转后两个点都要 pushup，且先更新下沉的那个。
  4. at(k) 的边界。k 从 1 开始，且必须满足 1≤k≤1 \le k \le 总数，否则会返回 0 号点。
  5. 多组数据忘记清空。root 和 idx 都要归零，否则上一组的数据会串进来。

十、 练习

11 参考资料

  • byvoid：《随机平衡二叉查找树 Treap 的分析与应用》
  • 本站旧版 Treap 笔记（内容已并入本文）

12 相关题解

题目 OJ 难度
[蓝桥杯 2024 省 B] 拔河 luogu 普及+/提高-
【模板】普通平衡树 luogu 提高+/省选-