解题思路
主席树的区间第 K 大的树上版本。
从根结点到每个结点都建立一个版本的线段树即可。查询时 a 到 b 路径时需要按照 $root[a] + root[b] - root[lca(a,b)] - root[fa[lca(a,b)]]$ 计算。
参考代码
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
using namespace std;
const int MAXN = 100001;
const int DEG = 20;
struct edge {
int to, next;
} e[MAXN<<1];
struct node {
int sum, l, r;
} hjt[MAXN*40];
int a[MAXN], sorted[MAXN], num;
bool is_son[MAXN]; // 是否为子结点,唯一一个不为 true 的即是树根
int fa[MAXN][DEG]; // fa[i][j]表示结点i的第2^j个祖先
int deg[MAXN]; // 深度数组
int root[MAXN], cnt;
int head[MAXN], tot;
void AddEdge(int u, int v) {
e[tot].to = v;
e[tot].next = head[u];
head[u] = tot++;
}
int GetIdx(int v) {
return lower_bound(sorted+1, sorted+1+num, v) - sorted;
}
void Init() {
tot = 0;
memset(head, -1, sizeof(head));
cnt = 0;
root[0] = 0;
}
inline int CreateNode(int sum, int l, int r) {
int idx = ++cnt;
hjt[idx].sum = sum;
hjt[idx].l = l;
hjt[idx].r = r;
return idx;
}
// 新建一棵线段树,只沿更新路径新建出较上个版本有修改的结点
// 调用参数
// root: 插入后新生成的线段树的根结点会赋值到 root 中存储
// pre_rt: 上一棵线段树的根
// pos: 本次要插入的数在线段树中的位置
// l, r: 递归参数。默认填写 1, num
void Insert(int &root, int pre_rt, int pos, int l, int r) {
root = CreateNode(hjt[pre_rt].sum+1, hjt[pre_rt].l, hjt[pre_rt].r);
if(l == r) return;
int m = (l+r) >> 1;
if(pos <= m)
Insert(hjt[root].l, hjt[pre_rt].l, pos, l, m);
else Insert(hjt[root].r, hjt[pre_rt].r, pos, m+1, r);
}
void BFS(int r) {
queue<int> q;
deg[r] = 0;
fa[r][0] = r;
q.push(r);
Insert(root[r], root[0], GetIdx(a[r]), 1, num);
while(!q.empty()) {
int f = q.front();
q.pop();
for(int i=1; i<DEG; ++i) {
fa[f][i] = fa[fa[f][i-1]][i-1];
}
for(int i=head[f]; ~i; i=e[i].next) {
int v = e[i].to;
if(v == fa[f][0]) continue;
deg[v] = deg[f] + 1;
fa[v][0] = f;
q.push(v);
Insert(root[v], root[f], GetIdx(a[v]), 1, num);
}
}
}
// 倍增 LCA
int LCA(int u, int v) {
if(deg[u] > deg[v]) swap(u,v);
int hu = deg[u], hv = deg[v];
int tu = u, tv = v;
for(int det=hv-hu, i=0; det ; det>>=1, ++i) {
if(det&1) tv = fa[tv][i];
}
if(tu == tv) return tu;
for(int i=DEG-1; i>=0; --i) {
if(fa[tu][i] == fa[tv][i]) continue;
tu = fa[tu][i];
tv = fa[tv][i];
}
return fa[tu][0];
}
// 本函数适用于查询树上路径 (u, v) 中的第 k 小。
// 调用参数
// tu, tv, ta, tfa: u 的根,v 的根,LCA(u, v) 的根,LCA(u, v) 的父结点的根
// k: 要查询区间第几小
// l, r: 递归参数。默认填写 1, num
int Query(int tu, int tv, int ta, int tfa, int k, int l, int r) {
if(l == r) return l;
int m = (l+r) >> 1;
int sum = hjt[hjt[tu].l].sum + hjt[hjt[tv].l].sum - hjt[hjt[ta].l].sum - hjt[hjt[tfa].l].sum;
if(k <= sum)
return Query(hjt[tu].l, hjt[tv].l, hjt[ta].l, hjt[tfa].l, k, l, m);
else return Query(hjt[tu].r, hjt[tv].r, hjt[ta].r, hjt[tfa].r, k-sum, m+1, r);
}
int main(int argc, char const *argv[]) {
int n, m, u, v, k;
while(~ scanf("%d %d", &n, &m)) {
Init();
for(int i=1; i<=n; ++i) {
scanf("%d", &a[i]);
sorted[i] = a[i];
}
sort(sorted+1, sorted+1+n);
num = unique(sorted+1, sorted+1+n) - (sorted+1);
for(int i=0; i<n-1; ++i) {
scanf("%d %d", &u, &v);
AddEdge(u, v);
AddEdge(v, u);
is_son[v] = true;
}
// 查找树的根结点
int r;
for(int i=1; i<=n; ++i) {
if(!is_son[i]) {
r = i;
break;
}
}
BFS(r);
fa[r][0] = 0; // 令根结点的父结点为 0,防止 Query 时重复计算导致错误
while(m--) {
scanf("%d %d %d", &u, &v, &k);
printf("%d\n", sorted[Query(root[u], root[v], root[LCA(u, v)], root[fa[LCA(u, v)][0]], k, 1, num)]);
}
}
return 0;
}
还不快抢沙发