Count on a tree

题目链接:Click here

Solution:

考虑主席树事实上是运用了一种前缀和的思想，每个点以它父亲为蓝本建树

那么答案为\(sz[u]+sz[v]-sz[lca(u,v)]-sz[fa(lca(u,v))]\)

Code:

#include<cstdio>
#include<cmath>
#include<queue>
#include<ctype.h>
#include<algorithm>
#include<iostream>
using namespace std;
const int N=2e5+1;
int n,m,cnt,head[N];
int num,tot,lstans,a[N],b[N];
int rt[N],sz[N*40],ls[N*40],rs[N*40];
int dep[N],f[35][N];
struct Edge{int nxt,to;}edge[N<<1];
void ins(int x,int y){
	edge[++cnt].nxt=head[x];
	edge[cnt].to=y;head[x]=cnt;
}
void dfs(int x,int fa){
	dep[x]=dep[fa]+1;f[0][x]=fa;
	for(int i=head[x];i;i=edge[i].nxt){
		int y=edge[i].to;
		if(!dep[y]) dfs(y,x);
    }
}
void trans(){int u=ceil(log2(n));
	for(int i=1;i<=u;i++)
		for(int j=1;j<=n;j++)
			f[i][j]=f[i-1][f[i-1][j]];
}
int lca(int x,int y){
	if(dep[x]<dep[y]) swap(x,y);
	if(x==y||y==1) return y;
	if(dep[x]!=dep[y]){
		for(int i=ceil(log2(n));i>=0;i--)
			if(dep[f[i][x]]>=dep[y]) x=f[i][x];
	}
	if(x==y) return x;
	for(int i=ceil(log2(n));i>=0;i--)
		if(f[i][x]!=f[i][y]) x=f[i][x],y=f[i][y];
	return f[0][x];
}
void build(int &q,int l,int r){
	q=++num;
	if(l==r) return ;
	int mid=l+r>>1;
	build(ls[q],l,mid);
	build(rs[q],mid+1,r);
}
void ins(int &q,int lst,int l,int r,int x){
	q=++num;sz[q]=sz[lst]+1;
	ls[q]=ls[lst],rs[q]=rs[lst];
	if(l==r) return ;
	int mid=l+r>>1;
	if(mid>=x) ins(ls[q],ls[lst],l,mid,x);
	else ins(rs[q],rs[lst],mid+1,r,x);
}
int query(int q,int p,int f1,int f2,int l,int r,int x){
	while(l!=r){int mid=l+r>>1;
		int v=sz[ls[q]]+sz[ls[p]]-(sz[ls[f1]]+sz[ls[f2]]);
		if(v>=x) q=ls[q],p=ls[p],f1=ls[f1],f2=ls[f2],r=mid;
		else q=rs[q],p=rs[p],f1=rs[f1],f2=rs[f2],l=mid+1,x-=v;
	}return b[l];
}
queue<int>q;
void makert(){
	q.push(1);
	while(!q.empty()){
		int x=q.front();q.pop();
		int v=lower_bound(b+1,b+tot+1,a[x])-b;
		ins(rt[x],rt[f[0][x]],1,tot,v);
		for(int i=head[x];i;i=edge[i].nxt){
			int y=edge[i].to;
			if(y!=f[0][x]) q.push(y);
		}
	}
}
int read(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)){if(ch=='-')f=-f;ch=getchar();}
	while(isdigit(ch)){x=x*10+ch-48;ch=getchar();}
	return x*f;
}
signed main(){
	n=read(),m=read();
	for(int i=1;i<=n;i++)
		a[i]=read(),b[++tot]=a[i];
	sort(b+1,b+tot+1);
	tot=unique(b+1,b+tot+1)-b-1;
	for(int i=1;i<n;i++){
		int x=read(),y=read();
		ins(x,y),ins(y,x);
	}dfs(1,0);trans();
	build(rt[0],1,tot);makert();
	for(int i=1;i<=m;i++){
		int u=read(),v=read(),k=read();
		u=u^lstans;int w=lca(u,v),fw=f[0][w];
		lstans=query(rt[u],rt[v],rt[w],rt[fw],1,tot,k);
		cout<<lstans<<endl;
	}
	return 0;
}