SPOJ - COT Count on a tree [主席树]

题意：给一颗树和树上每个点的权值，求树上节点u到v路径上的第k大。

题解：建立从根节点出发的主席树，当前节点的情况从其父亲节点update而来，计算路径上的第k大时，要结合如图所示的两个区间考虑。

AC代码:

#include<stdio.h>
#include<vector>
#include<algorithm>
#include<map>
#define N 100005
using namespace std;
int a[N],fa[N],b[N],f[N][20],dep[N];
int root[N*40],tree[N*40],lchild[N*40],rchild[N*40];
int tot,top;
vector<int>vt[N];
map<int,int>mp;
void update(int last,int cur,int x,int L,int R)
{
	tree[cur]=tree[last]+1;
	lchild[cur]=lchild[last];
	rchild[cur]=rchild[last];
	if(L==R)return ;
	int mid=L+R>>1;
	if(x<=mid)update(lchild[last],lchild[cur]=++tot,x,L,mid);
	else update(rchild[last],rchild[cur]=++tot,x,mid+1,R);
}
void dfs(int u,int ff,int deep)
{
	dep[u]=deep;
	update(root[ff],root[u]=++tot,mp[a[u]],1,top-1);
	//printf(">>%d %d\n",u,mp[a[u]]);
	f[u][0]=fa[u]=ff;
	for(int i=0;i<vt[u].size();i++)
	{
		int to=vt[u][i];
		if(to==ff)continue;
		dfs(to,u,deep+1);
	}
}
int lca(int x,int y){
    if(dep[x]<dep[y])swap(x,y);
    for(int i=16;i>=0;i--)if(dep[f[x][i]]>=dep[y])x=f[x][i];
    if(x==y)return x;
    for(int i=16;i>=0;i--)if(f[x][i]!=f[y][i])x=f[x][i],y=f[y][i];
    return f[x][0];
}
int query(int last1,int cur1,int last2,int cur2,int L,int R,int k)
{
	if(L==R)return L;
	int mid=L+R>>1;
	int flag=0;
	int lsum=tree[lchild[cur1]]-tree[lchild[last1]]+tree[lchild[cur2]]-tree[lchild[last2]];
	if(k<=lsum)return query(lchild[last1],lchild[cur1],lchild[last2],lchild[cur2],L,mid,k);
	else return query(rchild[last1],rchild[cur1],rchild[last2],rchild[cur2],mid+1,R,k-lsum);
}
int main()
{
	int n,m;
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		b[i]=a[i];
	}
	sort(b+1,b+n+1);
	top=unique(b+1,b+1+n)-b;
	for(int i=1;i<top;i++)
		mp[b[i]]=i;
	for(int i=0;i<n-1;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		vt[u].push_back(v);
		vt[v].push_back(u);
	}
	dfs(1,0,1);
	for(int k=1;k<=17;k++)
		for(int i=1;i<=n;i++)
			f[i][k]=f[f[i][k-1]][k-1];
	int kkk=0;
	int lc;
	for(int i=0;i<m;i++)
	{
		int u,v,k;
		scanf("%d%d%d",&u,&v,&k);
		lc=lca(u,v);
		kkk=b[query(root[f[lc][0]],root[u],root[lc],root[v],1,top-1,k)];
		printf("%d\n",kkk);
		
	}
}