SPOJ 10628 /LuoguSP10628 COT - Count on a tree

版权声明：原创文章，转载要注明作者哦 https://blog.csdn.net/DYT_B/article/details/83475118

题目描述：戳这里
题解：
这题是主席树上树（？？？）
我们发现对于一个点，我们可以维护从它到根的权值线段树，那么对于一个点对x，y，我们只需要求出它们的lca z,以及z的父节点，那么就可以容斥一下， $sumx+sumy-sumz-sum faz$。直接建n颗主席树是会mle的，所以我们直接用主席树，一个点的线段树建立在它的父节点上。

代码如下：

#pragma GCC optimize(2)
#pragma GCC optimize("Ofast")
#include<cstdio>
#include<string>
#include<algorithm>
#include<map>
using namespace std;
const int maxn=100005;
int n,m,tot,cnt,lnk[maxn],son[2*maxn],nxt[2*maxn];
int root[maxn],f[maxn][21],lson[18*maxn],rson[18*maxn],c[18*maxn],dep[maxn];
int a[maxn],b[maxn],rk[maxn];
map<int,int> id;
void add(int x,int y){
	son[++tot]=y,nxt[tot]=lnk[x],lnk[x]=tot;
}
void upd(int x,int y){lson[x]=lson[y],rson[x]=rson[y],c[x]=c[y];}
void insert(int &x,int y,int l,int r,int num){
	x=++tot; upd(x,y);
	if (l==r) {c[x]++; return;}
	int mid=(l+r)>>1;
	if (num<=mid) insert(lson[x],lson[y],l,mid,num);
	else insert(rson[x],rson[y],mid+1,r,num);
	c[x]=c[lson[x]]+c[rson[x]];
}
int query(int x,int y,int z,int z1,int l,int r,int kth){
	if (l==r) return rk[l];
	int mid=(l+r)>>1,num=c[lson[x]]+c[lson[y]]-c[lson[z]]-c[lson[z1]];
	if (kth<=num) return query(lson[x],lson[y],lson[z],lson[z1],l,mid,kth);
	else return query(rson[x],rson[y],rson[z],rson[z1],mid+1,r,kth-num);
}
void dfs(int x,int fa){
	insert(root[x],root[fa],1,cnt,id[a[x]]);
	f[x][0]=fa; dep[x]=dep[fa]+1;
	for (int i=1;i<=20;i++) f[x][i]=f[f[x][i-1]][i-1];
	for (int j=lnk[x];j;j=nxt[j])
		if (son[j]!=fa) dfs(son[j],x);
}
int lca(int x,int y){
	if (dep[x]<dep[y]) swap(x,y);
	for (int j=20;j>=0;j--) if (dep[f[x][j]]>=dep[y]) x=f[x][j];
	if (x==y) return x;
	for (int j=20;j>=0;j--) if (f[x][j]!=f[y][j]) x=f[x][j],y=f[y][j]; 
	return f[x][0];
}
int main(){
	scanf("%d%d",&n,&m);
	for (int i=1;i<=n;i++) scanf("%d",&a[i]),b[i]=a[i];
	sort(b+1,b+1+n);
	for (int i=1;i<=n;i++)
		if (!id[b[i]]) id[b[i]]=++cnt,rk[cnt]=b[i];
	for (int i=1;i<n;i++){
		int x,y; scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	tot=0; dfs(1,0);
	for (int i=1;i<=m;i++){
		int x,y,z,z1,k; scanf("%d%d%d",&x,&y,&k);
		z=lca(x,y); z1=f[z][0];
		printf("%d\n",query(root[x],root[y],root[z],root[z1],1,cnt,k));
	}
	return 0;
}