SPOJ - COT(主席树 LCA)

#include
#include
#include
#include
#include
using namespace std;
#define ll long long
const int N=1e5+10;
int h[N],e[N*2],ne[N*2],idx;
ll a[N];
vector<ll> v;
struct Node{
	int l,r,cnt;
}t[21*N];
int root[N],ans,deep[N],f[N][20];
int find(ll x)
{
	return lower_bound(v.begin() ,v.end() ,x)-v.begin() ;
}
void add(int a,int b)
{
	e[idx]=b;
	ne[idx]=h[a];
	h[a]=idx++;
}
int build(int l ,int r)
{
	int q=++ans;
	if(l==r)
	{
		return q;
	}
	int mid=(l+r)/2;
	t[q].l=build(l,mid);
	t[q].r=build(mid+1,r);
	return q;
}
int insert(int p,int l,int r,int x)
{
	int q=++ans;
	t[q]=t[p];
	if(l==r)
	{
		t[q].cnt++;
		return q;
	}
	int mid=(l+r)/2;
	if(x<=mid) t[q].l=insert(t[p].l,l,mid,x);
	else t[q].r=insert(t[p].r,mid+1,r,x);
	t[q].cnt=t[t[q].l].cnt+t[t[q].r].cnt;
	return q;
 } 
void dfs(int x,int fa)
{
	deep[x]=deep[fa]+1;
	root[x]=insert(root[fa],0,v.size()-1,find(a[x]));
	f[x][0]=fa;
	for(int i=1;i<=17;i++)
	{
		f[x][i]=f[f[x][i-1]][i-1];
	}
	for(int i=h[x];i!=-1;i=ne[i])
	{
		int j=e[i];
		if(j==fa) continue;
		dfs(j,x);
	}
}
int lca(int a,int b)
{
	if(deep[a]<deep[b]) return lca(b,a);
	for(int i=17;i>=0;i--)
	{
		if(deep[f[a][i]]>=deep[b])
		{
			a=f[a][i];
		}
	}
	if(a==b) return a;
	for(int i=17;i>=0;i--)
	{
		if(f[a][i]!=f[b][i])
		{
			a=f[a][i];
			b=f[b][i];
		}
	}
	return f[a][0];
}
int query(int u,int v,int lc,int flc ,int l,int r,int k)
{
	
	if(l==r)
	{
		return l;
	}
	int num=t[t[u].l].cnt+t[t[v].l].cnt-t[t[lc].l].cnt-t[t[flc].l].cnt;
	int mid=(l+r)/2;
	if(num>=k) return query(t[u].l,t[v].l,t[lc].l,t[flc].l,l,mid,k);
	else return query(t[u].r,t[v].r,t[lc].r,t[flc].r,mid+1,r,k-num);
}
int main()
{
	int n,m;
	cin>>n>>m;	
	memset(h,-1,sizeof h);
	for(int i=1;i<=n;i++)
	{
		scanf("%lld",&a[i]);
		v.push_back(a[i]);
	}
	sort(v.begin(),v.end());
	v.erase(unique(v.begin(),v.end()),v.end());
	for(int i=1;i<n;i++)
	{
		int a,b;
		scanf("%d%d",&a,&b);
		add(a,b);add(b,a);
	}
	root[0]=build(0,v.size()-1);
	deep[0]=0;
	dfs(1,0);
	for(int i=1;i<=m;i++)
	{
		int u,vv,k;
		scanf("%d%d%d",&u,&vv,&k);	
		printf("%lld\n",v[query(root[u],root[vv],root[lca(u,vv)],root[f[lca(u,vv)][0]],0,v.size()-1,k)]);	
	}
}