spoj Count on a tree【主席树】

题意：给一棵树还有一些询问求从u-v这条路径上点的权值的第K大。

我们可以从根开始维护主席树，就相当于以根为起点维护前缀和，这样我们可以求出u-v这条路径上对应点的权值，对u,v求一下lca(最近公共祖先)，然后点的权值就是c[u]+c[v]-2*[lca(u,v)]，然后同样二分找第k大。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;
const int MAXN=100010;
struct EDGE
{
	int v,next;
}edge[MAXN<<1];
int dep[MAXN<<1],rank[MAXN],num,set[MAXN<<1];
bool vis[MAXN];
int head[MAXN],size;
int Min(int x,int y)
{
	return dep[x]<dep[y]?x:y;
}
void init()
{
	memset(head,-1,sizeof(head));
	size=0;
}
void add_edge(int u,int v)
{
	edge[size].v=v;
	edge[size].next=head[u];
	head[u]=size++;
}
void dfs(int u,int deep)
{
    vis[u]=1;
    rank[u]=num;
    set[num]=u;
    dep[num++]=deep;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(vis[v])
          continue;
        dfs(v,deep+1);
        set[num]=u;
        dep[num++]=deep;
    }
}
int dp[MAXN<<1][21];
void RMQ_init()
{
    int i,j;
    for(i=0;i<num;i++)
      dp[i][0]=i;
    for(j=1;(1<<j)<num;j++)
    {
        for(i=0;i+(1<<j)<num;i++)
        {
            dp[i][j]=Min(dp[i][j-1],dp[i+(1<<(j-1))][j-1]);
        }
    }
}
int RMQ(int l,int r)
{
    if(r<l)
      swap(l,r);
    int k=(int)(log(r-l+1.0)/log(2.0));
    return Min(dp[l][k],dp[r-(1<<k)+1][k]);
}
int LCA(int l,int r)
{
    if(l>r)
      swap(l,r);
    int x=RMQ(l,r);
    return set[x];
}
int lson[MAXN*40],rson[MAXN*40],c[MAXN*40];
int tot,n,m;
int T[MAXN];
int vec[MAXN],idx;
int build(int l,int r)
{
	int root=tot++;
	c[root]=0;
	if(l!=r)
	{
		int mid=(l+r)>>1;
		lson[root]=build(l,mid);
		rson[root]=build(mid+1,r);
	}
	return root;
}
int update(int root,int pos,int val)
{
	int newroot=tot++;
	int tmp=newroot;
	c[newroot]=c[root]+val;
	int l=1,r=idx;
	while(l<r)
	{
		int mid=(l+r)>>1;
		if(pos<=mid)
		{
			lson[newroot]=tot++;
			rson[newroot]=rson[root];
			newroot=lson[newroot];
			root=lson[root];
			r=mid;
		}
		else
		{
			rson[newroot]=tot++;
			lson[newroot]=lson[root];
			newroot=rson[newroot];
			root=rson[root];
			l=mid+1;
		}
		c[newroot]=c[root]+val;
	}
    return tmp;
}
int query2(int l_root,int r_root,int fa_root,int pos,int k)
{
	int l=1,r=idx;
	while(l<r)
	{
		int mid=(l+r)>>1;
		int tmp=c[lson[l_root]]+c[lson[r_root]]-2*c[lson[fa_root]]+(pos>=l&&pos<=mid);
		if(k<=tmp)
		{
			l_root=lson[l_root];
			r_root=lson[r_root];
			fa_root=lson[fa_root];
			r=mid;
		}
		else
		{
			k-=tmp;
			l_root=rson[l_root];
			r_root=rson[r_root];
			fa_root=rson[fa_root];
			l=mid+1;
		}
	}
	return l;
}
int HASH(int val)
{
	int l=1,r=idx;
	while(l<=r)
	{
		int mid=(l+r)>>1;
		if(vec[mid]==val)
			return mid;
		if(vec[mid]<val)
			l=mid+1;
		else
			r=mid-1;
	}
	return l;
}
int val[MAXN];
void dfs1(int u,int fa)
{
	for(int i=head[u];i!=-1;i=edge[i].next)
	{
		int v=edge[i].v;
		if(v==fa)
			continue;
        int x=HASH(val[v]);
		T[v]=update(T[u],x,1);
		dfs1(v,u);
	}
}
void solve()
{
    tot=0;
	int i;
	for(i=1;i<=n;i++)
	{
		scanf("%d",&val[i]);
		vec[i]=val[i];
	}
	sort(vec+1,vec+1+n);
	idx=1;
	for(i=2;i<=n;i++)
	{
		if(vec[idx]!=vec[i])
			vec[++idx]=vec[i];
	}
	init();
	memset(vis,0,sizeof(vis));
	for(i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		add_edge(u,v);
		add_edge(v,u);
	}
    num=0;
	dfs(1,1);
	RMQ_init();
	T[n+1]=build(1,idx);
	T[1]=update(T[n+1],HASH(val[1]),1);
	dfs1(1,-1);
	while(m--)
	{
		int u,v,k;
		scanf("%d%d%d",&u,&v,&k);
		int fa=LCA(rank[u],rank[v]);
		int x=HASH(val[fa]);
		printf("%d\n",vec[query2(T[u],T[v],T[fa],x,k)]);
	}
}
int main()
{
	while(scanf("%d%d",&n,&m)==2)
		solve();
	return 0;
}
/*
8 5
105 2 9 3 8 5 7 7
1 2
1 3
1 4
3 5
3 6
3 7
4 8
*/