[SMOJ2165]treap_100865 treap]-优快云博客

本文介绍了一种基于Treap的数据结构实现，重点讲解了如何通过维护结点信息来快速查询第k小元素。提供了两种实现方式：懒删除和旋转删除，并附上了完整的参考代码。

Treap 模板题，插入和删除参见 Treap总笔记，这里简要谈谈如何在一棵 Treap 中求第 $k$ 小的值。

我们不妨在每一个结点 $u$ 记录附加值 $cnt$ 和 $size$ ，其中 $cnt$ 表示与 $u$ 的值重复的数的个数（这样就可以使用 lazy delete 的方法）， $size$ 表示以 $u$ 为根的子树的总数字个数，即

s i z e u = c n t u + \sum s i z e v for each child v of u

$size_u=cnt_u+\sum size_v\text{ for each child }v\text{ of }u$

为什么要维护这两个值呢？是为了方便我们后面的判断。

对于在当前以 $u$ 为根的子树，求第 $k$ 小的值时，不妨分类讨论答案的范围。记左子结点为 $l$ ，则：

当 $k\leq size_l$ 时，考虑到左子树的值均小于当前结点，则答案一定在左子树中，递归到左子树找第 $k$ 小
当 $size_l<k\leq size_l+cnt_u$ 时，要求的正好是当前重复的一段，当前结点即为答案，返回当前结点的值
不满足上述两种情况，答案在右子树，但因为当前子树中已有 $size_l+cnt_u$ 个数肯定小于答案，因此问题转化为求右子树第 $k-size_l-cnt_u$ 小

这样即可在 $O(log_2N)$ 的时间内完成询问。

需要注意的是，在代码中，需要在插入、删除和旋转后重新维护上述两个值（前两者是因为子树中对应的值已经发生了变化，旋转是因为子树的关系变化）。原则上必须按照自底向上的顺序维护，否则会出问题。因此旋转后，应先维护原根结点的值（现在它是子结点），再维护被旋上去的结点的值（现在它是根）。

参考代码一（懒删除）：

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

using namespace std;

struct Tnode {
    Tnode *child[2];
    int fix, val;
    int siz, cnt;
    Tnode (int v = 0) : fix(rand()), val(v), siz(1), cnt(1) { //构造函数
        child[0] = child[1] = 0;
    }

    void update() { //维护子树信息
        int lsiz = (child[0] ? child[0] -> siz : 0); //注意要子结点非空才能取值，否则读非法内存
        int rsiz = (child[1] ? child[1] -> siz : 0);
        siz = cnt + lsiz + rsiz;
    }
};

struct Treap {
    Tnode *root;
    Treap() { root = 0;/*new Tnode;*/ }

    void rotate(Tnode *&cur, int dir) { //dir=0左旋，1右旋
        Tnode *ch = cur -> child[dir ^ 1];
        cur -> child[dir ^ 1] = ch -> child[dir];
        ch -> child[dir] = cur;

        cur -> update(); //自下向上的维护顺序
        ch -> update();
        cur = ch;
    }

    void insert_val(Tnode *&cur, int v) {
        if (cur == NULL) cur = new Tnode(v);
        else {
            if (v == cur -> val) ++(cur -> cnt); //重复结点不新增
            else {
                int t = v > cur -> val;
                insert_val(cur -> child[t], v);
                if (cur -> child[t] -> fix < cur -> fix) rotate(cur, t ^ 1); //维护修正值的最小堆性质
            }
            cur -> update();
        }
    }

    int query_kth(Tnode *&cur, int k) {
        if (cur == NULL) return 0; //当前为空时及时返回 0，否则下面读非法内存
        if (k > cur -> siz) return 0; //要查的数排名比整个子树大小还大，无解
        int lsiz = (cur -> child[0] ? cur -> child[0] -> siz : 0);
        //分类讨论
        if (k <= lsiz) return query_kth(cur -> child[0], k);
        else if (k <= lsiz + cur -> cnt) return cur -> val;
        else return query_kth(cur -> child[1], k - lsiz - cur -> cnt);
    }

    void remove_val(Tnode *&cur, int v) {
        if (cur -> val == v) { //只改动标记
            -- cur -> cnt;
            -- cur -> siz;
            return;
        }
        remove_val(cur -> child[v > cur -> val], v); //直接修改标记就不需要旋转了
        cur -> update(); //删完维护一下信息
    }
} lkb_treap;

int main(void) {
    freopen("2165.in", "r", stdin);
    freopen("2165.out", "w", stdout);
    int N; scanf("%d", &N);
    while (N--) {
        int ch, k; scanf("%d%d", &ch, &k);
        if (ch == 1) lkb_treap.insert_val(lkb_treap.root, k);
        else if (ch == 2) printf("%d\n", lkb_treap.query_kth(lkb_treap.root, k));
        else lkb_treap.remove_val(lkb_treap.root, k);
    }
    return 0;
}

参考代码二（旋转删除）：

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

using namespace std;

struct Tnode {
    Tnode *child[2];
    int fix, val;
    int siz;
    Tnode (int v = 0) : fix(rand()), val(v), siz(1) {
        child[0] = child[1] = 0;
    }

    void update() {
        int lsiz = (child[0] ? child[0] -> siz : 0);
        int rsiz = (child[1] ? child[1] -> siz : 0);
        siz = lsiz + 1 + rsiz;
    }
};

struct Treap {
    Tnode *root;
    Treap() { root = 0;/*new Tnode;*/ }

    void rotate(Tnode *&cur, int dir) { //dir=0左旋，1右旋
        Tnode *ch = cur -> child[dir ^ 1];
        cur -> child[dir ^ 1] = ch -> child[dir];
        ch -> child[dir] = cur;

        cur -> update();
        ch -> update();
        cur = ch;
    }

    void insert_val(Tnode *&cur, int v) {
        if (cur == NULL) cur = new Tnode(v);
        else {
            int t = v > cur -> val; //插入到哪棵子树
            insert_val(cur -> child[t], v); //小于等于的都插左边，重复也新增
            if (cur -> child[t] -> fix < cur -> fix) rotate(cur, t ^ 1);
            cur -> update();
        }
    }

    int query_kth(Tnode *&cur, int k) {
        if (cur == NULL) return 0;
        if (k > cur -> siz) return 0;
        int lsiz = (cur -> child[0] ? cur -> child[0] -> siz : 0);
        if (k <= lsiz) return query_kth(cur -> child[0], k);
        else if (k == lsiz + 1) return cur -> val;
        else return query_kth(cur -> child[1], k - lsiz - 1);
    }

    void remove_val(Tnode *&cur, int v) {
        if (v == cur -> val) {
            //分类讨论
            if (!cur -> child[0] && !cur -> child[1]) cur = NULL; //叶子结点直接删
            else if (cur -> child[0] && cur -> child[1]) { //有两个儿子
                int t = cur -> child[0] -> fix < cur -> child[1] -> fix; //选修正值更小的代替自己以维护其最小堆性质
                rotate(cur, t);
                remove_val(cur -> child[t], v); //把旧的根旋下去之后递归删除
                cur -> update();
            } else if (cur -> child[0]) cur = cur -> child[0]; else cur = cur -> child[1]; //链结点直接代替
        } else {
            remove_val(cur -> child[v > cur -> val], v);
            cur -> update();
        }
    }
} lkb_treap;

int main(void) {
    freopen("2165.in", "r", stdin);
    freopen("2165.out", "w", stdout);
    int N; scanf("%d", &N);
    while (N--) {
        int ch, k; scanf("%d%d", &ch, &k);
        if (ch == 1) lkb_treap.insert_val(lkb_treap.root, k);
        else if (ch == 2) printf("%d\n", lkb_treap.query_kth(lkb_treap.root, k));
        else lkb_treap.remove_val(lkb_treap.root, k);
    }
    return 0;
}