带旋 Treap 解析

Treap

介绍

Treap 是一种依靠随机权值维护平衡性的平衡树，支持插入、删除、第 \(k\) 小、排名、前驱、后继等平衡树的基本功能。Treap 利用堆和平衡树的性质，其中堆使用的是随机优先级作为关键字。可以证明，平均单次复杂度是 \(O(\log n)\)（OI 不需要证明）。

算法内容

辅助函数和结构体

Node结构体

没啥好讲的。

struct Node {
    int lc, rc, rk, sz;
    int key;
} node[N];
int cnt;
int rt;

push_up和其他简化代码的函数

一样没啥好讲的。

int &lc(int u) {return node[u].lc;}
int &rc(int u) {return node[u].rc;}
int &rk(int u) {return node[u].rk;}
int &sz(int u) {return node[u].sz;}
int &key(int u) {return node[u].key;}
void push_up(int u) {
    sz(u) = sz(lc(u)) + sz(rc(u)) + 1;
}

旋转

Treap 的旋转分为左旋和右旋。

这是 \(1\) 的左旋。简单来说，就是 \(1\) 变成 \(1\) 的右孩子的左孩子， \(1\) 的右孩子原来的的左孩子变成 \(1\) 的右孩子。

void left(int &u) {
    int v = rc(u), tmp = lc(v);
    lc(v) = u;
    rc(u) = tmp;
    push_up(u), push_up(v);
    u = v;
}

右旋也一样，就是左旋反过来。

void right(int &u) {
    int v = lc(u), tmp = rc(v);
    rc(v) = u;
    lc(u) = tmp;
    push_up(u), push_up(v);
    u = v;
}

随机数

使用 mt19937，记得在构造里初始化种子。

mt19937 engine;
int randint(int minn = 0x80000000, int maxn = 0x7FFFFFFF) {
    std::uniform_int_distribution<int> dist(minn, maxn);
    return dist(engine);
}
Treap(int seed = random_device{}()) {
    rt = 0, cnt = 0;
    engine.seed(seed);
}

插入

由于平衡树性质更难维护，所以先依靠平衡树性质插入，再在回溯时进行旋转调整。

void insert(int &u, const int &ky) {
    if(!u) {
        node[++cnt] = {0, 0, randint(), 1, ky}, u = cnt;
        return;
    }
    sz(u)++;
    if(ky < key(u)) insert(lc(u), ky);
    else insert(rc(u), ky);
    if(lc(u) && rk(lc(u)) < rk(u)) right(u);
    if(rc(u) && rk(rc(u)) < rk(u)) left(u);
}

删除

删完再维护堆权值。

void erase(int &u, const int &ky) {
    if(!u) return;
    if(key(u) == ky) {
        if(!lc(u) || !rc(u)) u = lc(u) + rc(u);
        else if(rk(lc(u)) < rk(rc(u))) right(u), erase(rc(u), ky);
        else left(u), erase(lc(u), ky);
    } else if(ky < key(u)) erase(lc(u), ky);
    else erase(rc(u), ky);
    if(u) push_up(u);
}

排名、前驱、后继、第 \(k\) 小

跟普通的 BST 是一样的，不会的可以自己查。

int query_rk(int u, const int &ky) {
    if(!u) return 1;
    if(ky <= key(u)) return query_rk(lc(u), ky);
    else return sz(lc(u)) + 1 + query_rk(rc(u), ky);
}
int query_kth(int u, int k) {
    if(sz(lc(u)) >= k) return query_kth(lc(u), k);
    if(sz(lc(u)) + 1 == k) return key(u);
    return query_kth(rc(u), k - sz(lc(u)) - 1);
}
int pre(int u, const int &v) {
    if(!u) return -1e9;
    if(key(u) < v) return max(key(u), pre(rc(u), v));
    return pre(lc(u), v);
}
int suc(int u, const int &v) {
    if(!u) return 1e9;
    if(key(u) > v) return min(key(u), suc(lc(u), v));
    return suc(rc(u), v);
}

完整封装代码

const int N = 1e5 + 5;
class Treap {
private:
    mt19937 engine;
    struct Node {
        int lc, rc, rk, sz;
        int key;
    } node[N];
    int cnt;
    int rt;
    int &lc(int u) {return node[u].lc;}
    int &rc(int u) {return node[u].rc;}
    int &rk(int u) {return node[u].rk;}
    int &sz(int u) {return node[u].sz;}
    int &key(int u) {return node[u].key;}
    void push_up(int u) {
        sz(u) = sz(lc(u)) + sz(rc(u)) + 1;
    }
    int randint(int minn = 0x80000000, int maxn = 0x7FFFFFFF) {
        std::uniform_int_distribution<int> dist(minn, maxn);
        return dist(engine);
    }
    void right(int &u) {
        int v = lc(u), tmp = rc(v);
        rc(v) = u;
        lc(u) = tmp;
        push_up(u), push_up(v);
        u = v;
    }
    void left(int &u) {
        int v = rc(u), tmp = lc(v);
        lc(v) = u;
        rc(u) = tmp;
        push_up(u), push_up(v);
        u = v;
    }
    void insert(int &u, const int &ky) {
        if(!u) {
            node[++cnt] = {0, 0, randint(), 1, ky}, u = cnt;
            return;
        }
        sz(u)++;
        if(ky < key(u)) insert(lc(u), ky);
        else insert(rc(u), ky);
        if(lc(u) && rk(lc(u)) < rk(u)) right(u);
        if(rc(u) && rk(rc(u)) < rk(u)) left(u);
    }
    void erase(int &u, const int &ky) {
        if(!u) return;
        if(key(u) == ky) {
            if(!lc(u) || !rc(u)) u = lc(u) + rc(u);
            else if(rk(lc(u)) < rk(rc(u))) right(u), erase(rc(u), ky);
            else left(u), erase(lc(u), ky);
        } else if(ky < key(u)) erase(lc(u), ky);
        else erase(rc(u), ky);
        if(u) push_up(u);
    }
    int query_rk(int u, const int &ky) {
        if(!u) return 1;
        if(ky <= key(u)) return query_rk(lc(u), ky);
        else return sz(lc(u)) + 1 + query_rk(rc(u), ky);
    }
    int query_kth(int u, int k) {
        if(sz(lc(u)) >= k) return query_kth(lc(u), k);
        if(sz(lc(u)) + 1 == k) return key(u);
        return query_kth(rc(u), k - sz(lc(u)) - 1);
    }
    int pre(int u, const int &v) {
        if(!u) return -1e9;
        if(key(u) < v) return max(key(u), pre(rc(u), v));
        return pre(lc(u), v);
    }
    int suc(int u, const int &v) {
        if(!u) return 1e9;
        if(key(u) > v) return min(key(u), suc(lc(u), v));
        return suc(rc(u), v);
    }
public:
    Treap(int seed = random_device{}()) {
        rt = 0, cnt = 0;
        engine.seed(seed);
    }
    void insert(const int &ky) {
        insert(rt, ky);
    }
    void erase(const int &ky) {
        erase(rt, ky);
    }
    int query_rk(const int &ky) {
        return query_rk(rt, ky);
    }
    int query_kth(int k) {
        return query_kth(rt, k);
    }
    int pre(const int &ky) {
        return pre(rt, ky);
    }
    int suc(const int &ky) {
        return suc(rt, ky);
    }
};
Treap t;