[spoj COT - Count on a tree]树上第K小
分类：Data Structure Presidental Tree template
1. 题目链接
[spoj COT - Count on a tree]
2. 题意描述
N个节点的树，树上每个节点有一个点权。M次询问，每次询问一条链上的第
k
小数。
数据范围：(N,M<=100000) Time limit: 0.129s-0.516s Memory limit: 1536MB
3. 解题思路
本来以为跟COT2一样可以用莫队在树上跑。第k小的维护可以用pbds，线段树之类的东西维护。这样的复杂度是
O(nn√∗log2n)
。本题时间卡得比较紧，而且上述算法的复杂度偏高。 正解应该是主席树。看到给出的内存限制这么大应该联想到了... 在序列上，求第
k
小，步骤是先将点进行离散化，用主席树维护的是一个离散化之后点的前缀和。然后，查询的时候，对区间二分。
将情形推广到树上，同样的也是先将点进行离散化，然后从根节点，dfs向下依次将节点加入到主席树中。
那么询问(u→v)的路径上的点的时候，就是
sum[u]+sum[v]−2∗sum[lca(u,v)]+lca(u,v)
.
4. 实现代码
#include <bits/stdc++.h>
using namespace std;
#define FIN(x) freopen(#x".txt", "r", stdin)
const int MAXN =
110000;
const int DEEP =
20;
int n, m;
int w[MAXN], f[MAXN], fsz;
struct Edge {
int v, next;
} edge[MAXN <<
1];
int head[MAXN], tot, dep[MAXN], fa[MAXN][DEEP];
void init_edge() {
tot =
0;
memset(head, -
1,
sizeof(head));
}
inline void add_edge(
int u,
int v) {
edge[tot] = Edge{v, head[u]};
head[u] = tot ++;
}
int lca(
int u,
int v) {
while(dep[u] != dep[v]) {
if(dep[u] < dep[v]) swap(u, v);
int d = dep[u] - dep[v];
for(
int i =
0; i < DEEP; ++i) {
if(d >> i &
1) u = fa[u][i];
}
}
if(u == v)
return u;
for(
int i = DEEP -
1; i >=
0; --i) {
if(fa[u][i] != fa[v][i]) {
u = fa[u][i];
v = fa[v][i];
}
}
return fa[u][
0];
}
#define lch(u) nd[(u)].ch[0]
#define rch(u) nd[(u)].ch[1]
struct TNode {
int sum, ch[
2];
inline void reset() { sum = ch[
0] = ch[
1] =
0; }
} nd[MAXN *
20];
int ndIdx, root[MAXN];
void build(
int l,
int r,
int& rt) {
rt = ++ ndIdx;
nd[rt].reset();
if(l == r)
return;
int md = (l + r) >>
1;
build(l, md, lch(rt));
build(md +
1, r, rch(rt));
}
void update(
int p,
int l,
int r,
int& rt,
int prt) {
rt = ++ ndIdx;
nd[rt] = nd[prt];
++ nd[rt].sum;
if(l == r)
return;
int md = (l + r) >>
1;
if(p <= md) update(p, l, md, lch(rt), lch(prt));
else update(p, md +
1, r, rch(rt), rch(prt));
}
int kth(
int k,
int l,
int r,
int u,
int v,
int anc,
int ancp) {
if(l == r)
return l;
int md = (l + r) >>
1;
int sum = nd[lch(u)].sum + nd[lch(v)].sum - nd[lch(anc)].sum *
2 + (l <= ancp && ancp <= md);
if(sum >= k)
return kth(k, l, md, lch(u), lch(v), lch(anc), ancp);
else return kth(k - sum, md +
1, r, rch(u), rch(v), rch(anc), ancp);
}
int kth(
int k,
int u,
int v) {
int anc = lca(u, v), ancp = w[anc];
int le =
1, ri = fsz, md, sum;
u = root[u], v = root[v], anc = root[anc];
while(le < ri) {
md = (le + ri) >>
1;
sum = nd[lch(u)].sum + nd[lch(v)].sum - nd[lch(anc)].sum *
2 + (le <= ancp && ancp <= md);
if(sum >= k) {
ri = md;
u = lch(u), v = lch(v), anc = lch(anc);
}
else {
k -= sum;
le = md +
1;
u = rch(u), v = rch(v), anc = rch(anc);
}
}
return le;
}
void dfs(
int u,
int pre,
int deep) {
int v;
dep[u] = deep; fa[u][
0] = pre;
update(w[u],
1, fsz, root[u], root[pre]);
for(
int i = head[u]; ~i; i = edge[i].next) {
v = edge[i].v;
if(v == pre)
continue;
dfs(v, u, deep +
1);
}
}
int main() {
#ifdef ___LOCAL_WONZY___
FIN(input);
#endif
int u, v, k;
scanf(
"%d %d", &n, &m);
fsz =
0;
for(
int i =
1; i <= n; ++i) {
scanf(
"%d", &w[i]);
f[++ fsz] = w[i];
}
sort(f +
1, f + fsz +
1);
fsz = unique(f +
1, f + fsz +
1) - f -
1;
for(
int i =
1; i <= n; ++i) w[i] = lower_bound(f +
1, f + fsz +
1, w[i]) - f;
init_edge();
for(
int i =
2; i <= n; ++i) {
scanf(
"%d %d", &u, &v);
add_edge(u, v);
add_edge(v, u);
}
ndIdx =
0;
build(
1, fsz, root[
0]);
dfs(
1,
0,
0);
for(
int i =
1; i < DEEP; i++) {
for(
int j =
1; j <= n; j++) {
fa[j][i] = fa[fa[j][i -
1]][i -
1];
}
}
for(
int i =
1; i <= m; ++i) {
scanf(
"%d %d %d", &u, &v, &k);
int pos = kth(k, u, v);
printf(
"%d\n", f[pos]);
}
#ifdef ___LOCAL_WONZY___
cout <<
"Time elapsed: " <<
1.0 * clock() / CLOCKS_PER_SEC *
1000 <<
"ms." << endl;
#endif
return 0;
}
