AcWing 255. 第K小数-------主席树(可持久化线段树) 模板题

$AcWing$ $255$. 第$K$小数

一、题目大意

给定长度为 $N$ 的整数序列 $A$，下标为 $1$∼$N$

现在要执行 $M$ 次操作，其中第 $i$ 次操作为给出三个整数 $l_i,r_i,k_i$，求 $A[l_i],A[l_{i+1}],…,A[r_i]$ (即 $A$ 的下标区间 $[l_i,r_i]$)中第$k_i$ 小的数是多少

输入格式

第一行包含两个整数 $N$ 和 $M$

第二行包含 $N$ 个整数，表示整数序列$A$

接下来 $M$ 行，每行包含三个整数 $l_i,r_i,k_i$，描述第 $i$ 次操作

输出格式

对于每次操作输出一个结果，表示在该次操作中，第 $k$ 小的数的数值

每个结果占一行

数据范围

$N≤10^5,M≤10^4,|A[i]|≤10^9$

输入样例：

7 3
1 5 2 6 3 7 4
2 5 3
4 4 1
1 7 3

输出样例：

5
6
3

二、解题思路

主席树

据说主席树是一个叫 黄嘉泰 的人发明的，与我朝某位胡主席拼音简写同名，简写$hjt$，所以叫主席树。

一列数，可以对于每个点$i$都建一棵权值线段树，维护$1\sim i$这些数，每个不同的数出现的个数（权值线段树以值域作为区间）

现在，$n$棵线段树就建出来了，第$i$棵线段树代表$1\sim i$这个区间

例如，一列数，$n$为$6$，数分别为1 3 2 3 6 1
首先，每棵树都是这样的：

以第$4$棵线段树为例，$1\sim 4$的数分别为1 3 2 3

因为是同一个问题，$n$棵权值线段树的形状是一模一样的，只有节点的权值不一样
所以这样的两棵线段树之间是可以相加减的（两颗线段树相减就是每个节点对应相减）

想想，第$x$棵线段树减去第$y$棵线段树会发生什么？
第$x$棵线段树代表的区间是$[1,x]$
第$y$棵线段树代表的区间是$[1,y]$
两棵线段树一减
设$x>y，[1,x]−[1,y] = [y+1,x]$
所以这两棵线段树相减可以产生一个新的区间对应的线段树！

等等，这不是 前缀和的思想 吗?
这样一来，任意一个区间的线段树，都可以由我这$n$个基础区间表示出来了！
因为每个区间都有一个线段树
然后询问对应区间，在区间对应的线段树中查找$kth$就行了

这就是主席树的一个核心思想：前缀和思想

具体做法待会儿再讲，现在还有一个严峻的问题，就是$n$棵线段树空间太大了！
如何优化空间，就是主席树另一个核心思想

我们发现这$n$棵线段树中，有很多重复的点，这些重复的点浪费了大部分的空间，所以考虑如何去掉这些冗余点

在建树中优化

假设现在有一棵线段树，序列往右移一位，建一棵新的线段树
对于一个儿子的值域区间，如果权值有变化，那么新建一个节点，否则，连到原来的那个节点上

• 下面用加入一段序列来进行举例：序列$4$ $3$ $2$ $3$ $6$ $1$

区间$[1,1]$的线段树（蓝色节点为新节点）

当我们插入一个元素的时候，我们其实只需要修改其中一条链的数据，对于其他剩余的节点我们就直接复制，同时复制旧版本的节点到新的树上，复制完后新版本的节点加上新的版本号。

下面的执行步骤同理

区间$[1,2]$的线段树（橙色节点为新节点）

区间$[1,3]$的线段树（紫色节点为新节点）

这样是不是非常优秀啊？
主席树的思想就讲到这里，接下来具体的代码来实现它

由于主席树需要不停地开新节点，所以用完全二叉树的方式来存就没有必要了，因为除了一开始建的树以外，之后的版本里节点的左右孩子的下标都是不固定的。所以我们采用存指针的方式（即存两个孩子在数组中的下标），而对于当前节点维护的区间范围，可以作为参数在调用函数的时候传进来。

三、本题思路

首先要查找的是 某个下标区间里 的第$k$小数。回想一下平衡树里查找第$k$小数的过程，如果每个节点维护子树节点数，并且维护当前区间内数字出现次数，那么就可以通过类似折半查找的做法把第$k$小的数求出来。

本题也类似，可以建立一个线段树，每个节点维护的是$A$在该范围内的数的个数。例如线段树里维护$[0,2]$的区间的节点，记录的就是$A$中取值在$[0,2]$有多少个数。一开始版本$0$的线段树相当于在维护空数组，接着将$A[i]$逐次插入，形成$n$个版本。第$i$个版本维护的就是$A[1∼i]$在各个区间里取值的数的个数。如果要查询在$A[1]:A[r]$内的第$k$小的数，就可以查看第$r$个版本的线段树，然后每次看一下左孩子维护的区间里有多少个数，如果有$cnt$个，并且$k ≤ cnt$，那么就说明第$k$小的数在左半区间，则去左半区间找第$k$小的数；否则说明第$k$小的数在右半区间，则去右半区间找第$k − cnt$小的数。这和平衡树里求第$k$小的数的过程完全一样，也是在二分答案。

但是现在是要查询$A[l:r]$内第$k$小。这可以利用前缀和思想，我们考虑第$r$个版本和第$l − 1$个版本，两个版本的线段树的差，比如说比较两个版本维护区间$[a,b]$的节点里记录的$c$值，分别叫$c_r$ 和$c_{l-1}$，那么$c_r-c_{l-1}$其实就是$A[1:r]$相比于$A[1:l−1]$而言，在$[a,b]$里的数字个数多了多少个，那其实就是$A[l:r]$里有多少个数在$[a,b]$里。有了这个信息，就可以二分答案来解决了。本题由于$A[i]$的取值范围过大，按照这个取值范围建线段树太费空间，需要做 离散化，即将$A$映射到$0 ∼ n − 1$，然后用线段树维护$0 ∼ n− 1$这个区间即可。求完之后再映射回来即可。

四、主席树关键代码讲解

• 1、插入。这里的插入等价于普通线段树的单点修改，例如说插入$x$，对应的就是被维护数组的下标$x$的地方增加$1$（这里的$x$是离散化后的），只不过主席树会在每次插入的时候新开一个版本。

代码如下：

//经典的主席树插入
void insert(int &u, int l, int r, int x) {
    tr[++idx] = tr[u];  //新开一个节点idx++,将新节点指向旧的tr[u]
    tr[idx].cnt++;      //新节点的cnt，因为多插入了一个数字，所以个数+1,这样处理的话，省去了pushup
    u = idx;            //因为是地址引用，需要回写u等于idx

    if (l == r) return; //如果已经到了叶子节点，上面的操作就足够了，可以直接返回，否则需要继续向下递归

    int mid = l + r >> 1;
    if (x <= mid)
        insert(tr[u].l, l, mid, x); //因为tr[u]进入本函数时，最先把旧的复制过来，所以tr[u].l也是上一个版本的左儿子节点
    else
        insert(tr[u].r, mid + 1, r, x);
}

• 2、在$A[l:r]$里查询第$k$小。这需要在第$l − 1$版本和第$r$版本同时向下折半查找，每次都计算左半区间的元素个数$c$，然后和$k$比较，如果左半区间元素个数大于等于$k$ ，则说明答案在左半区间，去左半区间找第$k$小；否则说明在右半区间。去右半区间找第$k − c$小。

代码如下：

// p:前面的版本，q:后面的版本,[l,r]：控制的范围
// k:要查找第k小的数字
int query(int p, int q, int l, int r, int k) {
    if (l == r) return l;
    int mid = l + r >> 1;
    int cnt = tr[tr[q].l].cnt - tr[tr[p].l].cnt;
    if (k <= cnt)
        return query(tr[p].l, tr[q].l, l, mid, k);
    else
        return query(tr[p].r, tr[q].r, mid + 1, r, k - cnt);
}

五、静态数组离散化代码

#include <cstdio>
#include <algorithm>
using namespace std;
int read() {
    int x = 0, f = 1;
    char ch = getchar();
    while (ch < '0' || ch > '9') {
        if (ch == '-') f = -1;
        ch = getchar();
    }
    while (ch >= '0' && ch <= '9') {
        x = (x << 3) + (x << 1) + (ch ^ 48);
        ch = getchar();
    }
    return x * f;
}

const int N = 2e5 + 10;
int n, m;
int a[N], b[N], bl; // b和bl是一组，用于离散化的数组,bl为b的数组中有用数字的个数，一般下标0不放东西

struct Node {
    int l, r, cnt;
} tr[N << 5];
int root[N], idx;

//用于离散化的二分查找
int find(int x) {
    return lower_bound(b + 1, b + 1 + bl, x) - b;
}
//经典的主席树插入
void insert(int &u, int l, int r, int x) {
    tr[++idx] = tr[u];  //新开一个节点idx++,将新节点指向旧的tr[u]
    tr[idx].cnt++;      //新节点的cnt，因为多插入了一个数字，所以个数+1,这样处理的话，省去了pushup
    u = idx;            //因为是地址引用，需要回写u等于idx

    if (l == r) return; //如果已经到了叶子节点，上面的操作就足够了，可以直接返回，否则需要继续向下递归

    int mid = l + r >> 1;
    if (x <= mid)
        insert(tr[u].l, l, mid, x); //因为tr[u]进入本函数时，最先把旧的复制过来，所以tr[u].l也是上一个版本的左儿子节点
    else
        insert(tr[u].r, mid + 1, r, x);
}
// p:前面的版本，q:后面的版本,[l,r]：控制的范围
// k:要查找第k小的数字
int query(int p, int q, int l, int r, int k) {
    if (l == r) return l;
    int mid = l + r >> 1;
    int cnt = tr[tr[q].l].cnt - tr[tr[p].l].cnt;
    if (k <= cnt)
        return query(tr[p].l, tr[q].l, l, mid, k);
    else
        return query(tr[p].r, tr[q].r, mid + 1, r, k - cnt);
}

int main() {
//文件输入输出
#ifndef ONLINE_JUDGE
    freopen("P3834.in", "r", stdin);
#endif
    n = read(), m = read();
    for (int i = 1; i <= n; i++)
        a[i] = b[i] = read();

    //数据范围太大，直接建线段树会MLE，但是比较稀疏，可以离散化后用相对应的序号，数据量就小了
    sort(b + 1, b + 1 + n);
    bl = unique(b + 1, b + 1 + n) - b - 1; //离散化后共m个数字

    // 0号版本没有内容时，主席树是不需要进行build的,强行build时，可能会有部分测试点TLE
    // 0号版本有内容时，主席树是需要build的，不build，初始值无法给上

    // 主席树的数字增加，每增加一个，就相当于增加了一个版本root[i]记录了版本i的根节点
    for (int i = 1; i <= n; i++) {
        root[i] = root[i - 1];              //开新版本号i,抄袭上一个版本i-1的根节点
        insert(root[i], 1, bl, find(a[i])); //向版本i中增加find(a[i])的值
    }

    while (m--) {
        int l, r, k;
        l = read(), r = read(), k = read();
        //采用类似于前缀的方法，对位相减，由于是动态开点，需要指明控制范围[1,bl]
        //要查询的数字是k
        printf("%d\n", b[query(root[l - 1], root[r], 1, bl, k)]);
    }

    return 0;
}

六、$vector$离散化代码

#include <algorithm>
#include <cstdio>
#include <vector>
using namespace std;
//快读
int read() {
    int x = 0, f = 1;
    char ch = getchar();
    while (ch < '0' || ch > '9') {
        if (ch == '-') f = -1;
        ch = getchar();
    }
    while (ch >= '0' && ch <= '9') {
        x = (x << 3) + (x << 1) + (ch ^ 48);
        ch = getchar();
    }
    return x * f;
}
const int N = 2e5 + 10;
int n, m;
int a[N];
vector<int> ys;

struct Node {
    int l, r, cnt;
} tr[N << 5];
int root[N], idx;

// 0号版本没有内容时，主席树是不需要进行build的,强行build时，可能会有部分测试点TLE
// 0号版本有内容时，主席树是需要build的，不build，初始值无法给上

int find(int x) {
    return lower_bound(ys.begin(), ys.end(), x) - ys.begin();
}

//经典的主席树插入
void insert(int &u, int l, int r, int x) {
    tr[++idx] = tr[u];  //新开一个节点idx++,将新节点指向旧的tr[u]
    u = idx;            //因为是引用，为了回传正确值，需要u=idx-1
    tr[u].cnt++;        //新节点的cnt，因为多插入了一个数字，所以个数+1,这样处理的话，省去了pushup
    if (l == r) return; //如果已经到了叶子节点，上面的操作就足够了，可以直接返回，否则需要继续向下递归

    int mid = l + r >> 1;
    if (x <= mid)
        insert(tr[u].l, l, mid, x); //因为tr[u]进入本函数时，最先把旧的复制过来，所以tr[u].l也是上一个版本的左儿子节点
    else
        insert(tr[u].r, mid + 1, r, x);
}

int query(int p, int q, int l, int r, int k) {
    if (l == r) return l;
    int mid = l + r >> 1;
    int cnt = tr[tr[q].l].cnt - tr[tr[p].l].cnt;
    if (k <= cnt)
        return query(tr[p].l, tr[q].l, l, mid, k);
    else
        return query(tr[p].r, tr[q].r, mid + 1, r, k - cnt);
}

int main() {
//文件输入输出
#ifndef ONLINE_JUDGE
    freopen("P3834.in", "r", stdin);
#endif
    n = read(), m = read();
    for (int i = 1; i <= n; i++) {
        a[i] = read();
        ys.push_back(a[i]);
    }
    //数据范围太大，直接建线段树会MLE，但是比较稀疏，可以离散化后用相对应的序号，数据量就小了
    sort(ys.begin(), ys.end());
    ys.erase(unique(ys.begin(), ys.end()), ys.end());
   
    //主席树的数字增加，每增加一个，就相当于增加了一个版本root[i]记录了版本i的根节点
    for (int i = 1; i <= n; i++) {
        root[i] = root[i - 1];
        insert(root[i], 0, ys.size() - 1, find(a[i]));
    }

    while (m--) {
        int l, r, k;
        l = read(), r = read(), k = read();
        printf("%d\n", ys[query(root[l - 1], root[r], 0, ys.size() - 1, k)]);
    }

    return 0;
}

七、时间复杂度

时间复杂度$O(nlogn+mlogn)$，即每次插入和查询时间复杂度都是$O(logn)$，空间$O(n+nlogn)$，每次插入都要新开$O(logn)$个节点。