SPOJ Count on a tree（树上第K大）

题目

主席树+LCA

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
const int MAXN=(int)1e5+5;
struct edge{
	int to,nxt;
}ed[MAXN<<1];
int head[MAXN],cnt;
void addedge(int u,int v){
	ed[cnt].to=v;
	ed[cnt].nxt=head[u];
	head[u]=cnt++;
}
int root[MAXN],ls[MAXN*40],rs[MAXN*40],sum[MAXN*40],tot;
int a[MAXN],dep[MAXN],f[MAXN][20];
vector<int>ve;int n,m,vsz;
void insert(int l,int r,int &x,int y,int pos){
	x=++tot;ls[x]=ls[y],rs[x]=rs[y],sum[x]=sum[y]+1;
	if(l==r)return;
	int mid=l+r>>1;
	if(mid>=pos)insert(l,mid,ls[x],ls[y],pos);
	else insert(mid+1,r,rs[x],rs[y],pos);
}
int query(int l,int r,int u,int v,int a,int b,int k){
	if(l==r)return l;
	int mid=l+r>>1;
	int step=sum[ls[u]]+sum[ls[v]]-sum[ls[a]]-sum[ls[b]];
	if(step>=k)return query(l,mid,ls[u],ls[v],ls[a],ls[b],k);
	else return query(mid+1,r,rs[u],rs[v],rs[a],rs[b],k-step);
}

void dfs(int u,int pre){
	insert(1,vsz,root[u],root[pre],a[u]);
	for(int i=head[u];i!=-1;i=ed[i].nxt){ 
		int v=ed[i].to; 
		if(!dep[v]){			
		  dep[v]=dep[u]+1;
		  f[v][0]=u; 
		  dfs(v,u);
		}
	}
}
int lca(int a,int b){
	int t,i;
	if(dep[a]<dep[b])t=a,a=b,b=t;
	for(i=16;i>=0;i--){
		if(dep[f[a][i]]>=dep[b])a=f[a][i];
	}
	if(a==b)return a;
	for(i=16;i>=0;i--){
		if(f[a][i]!=f[b][i]&&f[a][i]!=0&&f[b][i]!=0)a=f[a][i],b=f[b][i];
	}
	return f[a][0];
}
int query(int u,int v,int k){
	int LCA=lca(u,v);
	return query(1,vsz,root[u],root[v],root[LCA],root[f[LCA][0]],k);
}
int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++)scanf("%d",&a[i]),ve.push_back(a[i]);
	sort(ve.begin(),ve.end());
	ve.erase(unique(ve.begin(),ve.end()),ve.end());
	for(int i=1;i<=n;i++)a[i]=lower_bound(ve.begin(),ve.end(),a[i])-ve.begin()+1;
	vsz=ve.size();
	for(int i=1;i<n;i++){
		int u,v;	
		scanf("%d%d",&u,&v);
		addedge(u,v);
		addedge(v,u);
	}
	dep[1]=1;
	dfs(1,0);
	for(int j=1;(1<<j)<=n;j++)for(int i=1;i<=n;i++)if(f[i][j-1]!=0)f[i][j]=f[f[i][j-1]][j-1];
    while(m--){
    	int u,v,k;
    	scanf("%d%d%d",&u,&v,&k);
    	printf("%d\n",ve[query(u,v,k)-1]);
	}
	return 0;
}