[算法学习]——可持久化权值线段树求区间第k小（主席树）

前置知识

树状数组，线段树 （本文不讲）

权值线段树，树上二分，动态开点线段树（下文会讲）

前置知识讲解

权值线段树

虽然叫线段树，但事实上我们一般用树状数组去写权值线段树

并非像线段树那样维护数组下标区间$[l,r]$中的信息，而是以权值为下标，维护权值为$[l,r]$的数量

一个简单的例子：询问数组中有多少满足$5\le x\le 8$的权值$x$

首先我们考虑如何维护权值线段树，对于第$i$个元素$a_i$，我们应该$bit.modify(a_i,1)$，表示权值为$a_i$的数量$+1$

考虑查询，树状数组是类似于前缀和的思想，$bit.query(x)$表示权值最多为$x$的数量，那么就是$bit.query(8)-bit.query(4)$，值最多为$8$的数量减去值最多为$4$的数量，那么我们就得到了权值处于区间$[5,8]$的数量

我们可以画个图简单的理解下：

在这个图中，表示数组有两个权值为一的元素，一个权值为二的元素，零个权值为四的元素，然后我们利用树状数组查询$bit.query(4)$相当于把所有小于等于$4$的值的数量都求出来了，即$bit.query(4)=2+1+2+0=5$

注：下文的主席树还是以用线段树去实现权值线段树，这里主要是讲思想，所以用了树状数组

树上二分

学会了权值线段树后，我们考虑求全局第$k$小，我们二分一个答案$mid$，然后计算权值处于区间$[l,mid]$的数量$cnt$。

如果$cnt<k$的话，那么显然答案不会在$[l,mid]$区间上，所以我们要判断区间$[mid+1,r]$中哪个数有可能是答案，此外，因为$[l,mid]$都不是答案，所以$k$要减去这些数量，即$k=k-cnt$

如果$cnt\ge k$的话，说明答案在$[l,mid]$区间内，那么我们继续迭代判断即可。

我们可以手动模拟一下：

例如我们在这个权值线段树上求全局第$3$大

一开始$l=1,r=6,mid=\frac{l+r}{2}=3$

然后$cnt=bit.query(l,mid)=5>k$

所以答案不可能在区间$[mid+1,r]=[4,6]$上

接下来$r=mid=3$，那么有$l=1,r=3,mid=2$

就有$cnt=bit.query(l,mid)=3>k$

那么$r=mid=2$，有$l=1,r=2,mid=1$

$cnt=bit.query(l,mid)=2<k$

就说明答案不会在$[l,mid]=[1,1]$区间中，因为$[1,1]$中所有的数都小于我们的答案，我们后面又不会考虑这部分的贡献，所以要$k-cnt=1$

最后看区间$l=2,r=2,mid=2$，这里显然可以得到答案是$2$了

注：线段树上二分也类似，不过是把区间查询操作修改成线段树的查询格式

动态开点

在普通线段树中，表示左右儿子的编号我们一般是当前编号$id$来表示：左儿子$id\ast 2$，右儿子$id\ast 2+1$

但是在主席树中，我们的每个节点的儿子有可能是新建节点，也有可能是依附于之前的线段树的某个节点。

所以我们选择给每个节点单独编号。

具体实现就是我们单独开一个结构体用来维护三个值——当前节点的信息，左儿子的编号，右儿子的编号

正文

核心思想

主席树的本质是在每个位置维护一个从初始空树到当前位的前缀和权值线段树

暴力思想

对于区间求第$k$小，首先我们现象一个非常暴力的做法：

我们先建一个值为空，但是保有线段树结构的权值线段树，这是我们的第$0$版本线段树

对一个长度为$n$的数组，对于第一个元素，我们先复制好刚刚建立的只有树的结构的权值线段树，一次建立一个新的权值线段树，然后把第一个元素对应的权值的节点进行加一操作，并且更新该权值线段树，并且把这个线段树成为第$1$版本线段树

对于第二个元素，我们把第$1$棵权值线段树复制过来再建一颗新的权值线段树，然后根据第二个元素的权值去更新该线段树，这就是我们的第$2$版本线段树

对于第$i$个元素，我们把第$i-1$版本线段树复制过来建新树，并且根据第$i$个元素对应的权值去修改新树，作为第$i$版本线段树

这样子我们有了$n+1$个版本的权值线段树，那么如果我们要求区间$[l,r]$的第$k$小该怎么做呢？

我们注意到第$i$版本权值线段树保留了第$i-1$版本线段树的全部信息（因为是复制过来的），并且维护了第$i$个元素的信息，所以这是一个类似于前缀和的信息维护

我们找到第$r$版本权值线段树和第$l-1$版本线段树，然后第$r$版本线段树的每个节点都减去第$l-1$版本线段树的对应节点的信息（因为结构都是根据最开始的空树确定的，所以结构都一样），这样子我们就得到了只包含区间$[l,r]$的元素的权值线段树（前缀和思想）。那么求区间第$k$小我们只需要在这棵权值线段树上进行树上二分就可以了，怎么二分这点在上文中讲过

优化思想

我们可以注意到每次维护第$i$个元素所代表的权值，在权值线段树上只修改了$\log n$个节点（单点修改），但是在暴力思想中我们却每一次都把所有节点都复制了一遍，这显然是非常不优的，那么我们有没有什么办法只维护修改的节点呢？

这个办法就是主席树的思想了。我们注意到因为只修改了$\log n$个节点，其余节点都和前一个版本的线段树是一样的，所以我们干脆让相同的节点作为共享节点，修改的$\log n$个节点我们都当成新开的节点，而没有修改的节点我们就直接拿新开的节点去连边，这样就相当于建了一棵新树（具体会在下文讲解）

但是我们该如何实现连边操作呢？这就需要动态开点了（见上文），我们对每个节点维护三个信息——【题目需要的信息，左儿子节点编号，右儿子节点编号】。这样子我们就可以知道当前节点的儿子的编号都是哪些节点了（具体会在下文讲解）

对于查询操作，类似于暴力，但是现在只有一棵树+新家的若干个新节点了，但是我们依旧是找到$r$和$l-1$版本对应的节点，然后将节点中与题目有关的信息相减即可（具体会在下文讲解）

建树

这里我们以维护一个权值范围为$[1,8]$的主席树为例

黑色部分是一颗值全部为零的权值线段树，我们只保留其作为完全二叉树的结构，这就是我们的第$0$版本线段树

同时每个节点要维护三个信息：题目所需的信息，左儿子编号，右儿子编号

struct node {
    int l, r, val;
    // 左儿子编号，右儿子编号，题目所需的信息，这里是元素数量
} seg[N];

建树我们应该改成动态开点的方式，因为之后会有节点是其他很多节点的共同的儿子

void build(int version, int l, int r) {
    if (l == r) {
        seg[version].val = 0;
        //初始为0
        return ;
    }
    int mid = (l + r) >> 1;
    seg[version].l = ++ id, seg[version].r = ++ id;
    //给左右儿子赋上编号
    build(seg[version].l, l, mid);
    build(seg[version].r, mid + 1, r);

    seg[version].val = seg[seg[version].l].val + seg[seg[version].r].val;
    //线段树的信息合并操作
}

修改

接下来我们维护第一个元素，其权值为$8$

那么我们首先先建立一个新的节点（编号$16$），作为我们第$1$版本线段树的根节点，然后我们发现权值$8$应该在右儿子上，所以左儿子不需要新建，那么我们就让上一个版本的线段树的左儿子作为当前节点的左儿子即可，就是$16$向$2$连一个边（在代码中，就是让节点的信息中的左儿子编号等于$2$），因为要修改的节点在右儿子，所以我们新建一个节点（编号17），并且连边（节点$16$的右儿子编号改为$17$）

以此类推，编号$18$节点左儿子不修改，连接上一版本的线段树的节点，右儿子新建节点。此时编号$19$的节点是叶子节点，其表示权值为$8$的元素的数量，所以让该节点的对问题所需的信息$+1$（这里“对问题所需的信息”是指求区间第$k$小的话我们需要维护权值的元素数量，所以这里指权值为$8$的元素数量加一），并像线段树一样返回去维护父亲即祖先节点的信息即可

注意，如果我们只看当前根节点所连接的节点，那么这就是暴力思想中的第一版本线段树，类似于下图

继续维护第二个元素，其权值为$4$

根据上面的做法，我们会得到这样的线段树：

我们依旧是先建一个新的节点作为新版本线段树的根节点（编号$20$），因为权值为$4$的节点在左子树，所以右儿子不需要改变，那么链接上一版本的对应节点的右儿子即可（即新根节点的右儿子是上一版本根节点的右儿子）

同理，把跟该根节点无关的节点去掉后就是上文暴力中的第二版本线段树了

void update(int l, int r, int last, int now, int pos, int val) {
    // 区间左右端点，上一个版本的对应节点的编号，当前节点的编号，需要求改的位置，修改的数量
    if (l == r) {
        seg[now].val = seg[last].val + val;
        // 在上一个版本对应节点的基础上进行操作
        return ;
    }
    int mid = (l + r) >> 1;
    seg[now].l = seg[last].l, seg[now].r = seg[last].r;
    // 先假设两个节点都不需要改
    if (pos <= mid) {
        // 如果要修改的点在左子树
        seg[now].l = ++ id;
        // 新建左节点
        update(l, mid, seg[last].l, seg[now].l, pos, val);
    }
    else {
        // 否则在右子树
        seg[now].r = ++ id;
        update(mid + 1, r, seg[last].r, seg[now].r, pos, val);
    }

    seg[now].val = seg[seg[now].l].val + seg[seg[now].r].val;
    // 线段树的信息合并操作
}

查询

首先简单回顾下权值线段树的每个节点表示什么

$1$号节点表示权值范围为$[1,8]$的元素数量

$6$号节点表示权值范围为$[5,6]$的元素数量

$9$号节点表示权值为$2$的元素数量

我们要求区间$[l,r]$中的第$k$小，那么首先我们要找到两个版本的线段树$last=l-1,now=r$，然后从根节点开始遍历，用两个参数表示两个版本线段树的对应节点编号，一步步的树上二分找到我们所需的答案即可

int query(int l, int r, int last, int now, int k) {
    if (l == r) return l;

    int mid = (l + r) >> 1, d = seg[seg[now].l].val - seg[seg[last].l].val;
    // d是版本[last, now]中左子树代表的范围内的元素数量
    if (d < k)
        return query(mid + 1, r, seg[last].r, seg[now].r, k - d);
    else
        return query(l, mid, seg[last].l, seg[now].l, k);
    // 树上二分部分，如果左子树代表范围的元素数量要小于k，说明答案在右子树，并且k要减去d，否则答案在左子树
}

全代码

例题：牛客周赛62-F-小红的中位数查询（hard）

#include <bits/stdc++.h>

using i64 = long long;
using u64 = unsigned long long;

typedef std::pair<long long, long long> PII;
const int mod = 998244353;
const int N = 2e6 + 1000;
const int INF = 0x3f3f3f3f;
const long long LINF = 1e18;
const double eps = 1e-12;
std::mt19937_64 rnd(std::random_device{}());

int n, m, root[N], id = 1;

struct node {
    int l, r, val;
} seg[N];

void build(int version, int l, int r) {
    if (l == r) {
        seg[version].val = 0;
        return ;
    }
    int mid = (l + r) >> 1;
    seg[version].l = ++ id, seg[version].r = ++ id;
    build(seg[version].l, l, mid);
    build(seg[version].r, mid + 1, r);

    seg[version].val = seg[seg[version].l].val + seg[seg[version].r].val;
}

void update(int l, int r, int last, int now, int pos, int val) {
    if (l == r) {
        seg[now].val = seg[last].val + val;
        return ;
    }
    int mid = (l + r) >> 1;
    seg[now].l = seg[last].l, seg[now].r = seg[last].r;
    if (pos <= mid) {
        seg[now].l = ++ id;
        update(l, mid, seg[last].l, seg[now].l, pos, val);
    }
    else {
        seg[now].r = ++ id;
        update(mid + 1, r, seg[last].r, seg[now].r, pos, val);
    }

    seg[now].val = seg[seg[now].l].val + seg[seg[now].r].val;
}

int query(int l, int r, int last, int now, int k) {
    if (l == r) return l;

    int mid = (l + r) >> 1, d = seg[seg[now].l].val - seg[seg[last].l].val;
    if (d < k)
        return query(mid + 1, r, seg[last].r, seg[now].r, k - d);
    else
        return query(l, mid, seg[last].l, seg[now].l, k);
}

void solve() {
    std::cin >> n >> m;
    std::vector<int> a(n + 1), b;
    for (int i = 1; i <= n; i ++) {
        std::cin >> a[i];
        b.emplace_back(a[i]);
    }
    std::sort(b.begin(), b.end());
    b.erase(std::unique(b.begin(), b.end()), b.end());

    int len = b.size();
    root[0] = 1;
    build(root[0], 1, len);

    auto get = [&](int x) {
        return std::lower_bound(b.begin(), b.end(), x) - b.begin() + 1;
    };

    for (int i = 1; i <= n; i ++) {
        int p = get(a[i]);
        root[i] = ++ id;
        update(1, len, root[i - 1], root[i], p, 1);
    }

    while (m --) {
        int l, r, k;
        std::cin >> l >> r;
        k = (r - l + 1 + 1) / 2;
        std::cout << b[query(1, len, root[l - 1], root[r], k) - 1] << "\n";
    }
}

signed main() {
    std::ios::sync_with_stdio(false), std::cin.tie(nullptr), std::cout.tie(nullptr);

    int tmp = 1;

    // std::cin >> tmp;

    while (tmp--)
        solve();

    return 0;
}

例题

P3834 【模板】可持久化线段树 2