Count on a tree

树上主席树练习

luogu P2633

#include<iostream>
#include<cstdio>
#include<algorithm>
#define N 100000
#define NlogN 2500000 
using namespace std;
int n,m,a[N+1],b[N+1],tree[N+1],l[NlogN],r[NlogN],sum[NlogN],deep[N+1],f[N+1][31],tot,len,mid,u,v,k,last;
int cnt,he[N<<2],nxt[N<<2],to[N<<2];
void add(int u,int v)
{
	cnt++;
	nxt[cnt]=he[u];
	he[u]=cnt;
	to[cnt]=v;
}
int build(int ll,int rr)
{
	int rt=++tot;
	if(ll<rr)
	{
		mid=(ll+rr)>>1;
		l[rt]=build(ll,mid);
		r[rt]=build(mid+1,rr);
	}
	return rt;
}
int update(int pre,int ll,int rr,int x)
{
	int rt=++tot;
	l[rt]=l[pre],r[rt]=r[pre],sum[rt]=sum[pre]+1;
	if(ll<rr)
	{
	  mid=(ll+rr)>>1;
	  if(x<=mid) l[rt]=update(l[pre],ll,mid,x);
	  else r[rt]=update(r[pre],mid+1,rr,x);	
	}
	return rt;
}
void build_tree(int pre,int now)
{
	deep[now]=deep[pre]+1,f[now][0]=pre;
	int rank=lower_bound(b+1,b+len+1,a[now])-b;
	tree[now]=update(tree[pre],1,len,rank);
	for(int i=he[now];i;i=nxt[i])
	{
		int g=to[i];
		if(g==pre) continue;
		build_tree(now,g);
	}
}
int lca(int q,int p)
{
	if(deep[q]>deep[p]) swap(q,p);
	int cz=deep[p]-deep[q],js=0;
	while(cz)
	{
		if(cz&1) p=f[p][js];
		js++,cz>>=1;
	}
	if(q==p) return q;
	for(int i=20;i>=0;i--) if(f[q][i]!=f[p][i]) q=f[q][i],p=f[p][i];
	return f[q][0];
}
int query(int t1,int t2,int t3,int t4,int ll,int rr,int id)
{
	if(ll==rr) return ll;
	int cz=sum[l[t1]]+sum[l[t2]]-sum[l[t3]]-sum[l[t4]];
	mid=(ll+rr)>>1;
	if(cz>=id) return query(l[t1],l[t2],l[t3],l[t4],ll,mid,id);
	else return query(r[t1],r[t2],r[t3],r[t4],mid+1,rr,id-cz);
}
int qread()
{
	char c=getchar();
	int x=0;
	while(c>'9'||c<'0') c=getchar();
	while(c<='9'&&c>='0') x=(x<<3)+(x<<1)+c-'0',c=getchar();
	return x;
}
void ks(int xx)
{
	if(xx>=10) ks(xx/10);
	putchar(xx%10+'0');
}
int main()
{
	n=qread(),m=qread();
	for(int i=1;i<=n;i++) a[i]=qread(),b[i]=a[i];
	sort(b+1,b+n+1);
	len=unique(b+1,b+n+1)-b-1;
	tree[0]=build(1,len);//建零号主席树 
	for(int i=1,x,y;i<=n-1;i++) x=qread(),y=qread(),add(x,y),add(y,x);
	build_tree(0,1); 
	for(int i=1;i<=20;i++) //ST表 
	for(int j=1;j<=n;j++) f[j][i]=f[f[j][i-1]][i-1];
	for(int i=1;i<=m;i++)
	{
		scanf("%d%d%d",&u,&v,&k);
		last=u^last;
		int LCA=lca(last,v);
		last=b[query(tree[last],tree[v],tree[LCA],tree[f[LCA][0]],1,len,k)];
		ks(last),putchar('\n');
	}
	return 0;
}