主席树+LCA+spoj10628

10628. Count on a tree Problem code: COT

You are given a tree with N nodes.The tree nodes are numbered from 1 to N.Each node has an integer weight.

We will ask you to perform the following operation:

• u v k : ask for the kth minimum weight on the path from node u to node v

 

Input

In the first line there are two integers N and M.(N,M<=100000)

In the second line there are N integers.The ith integer denotes the weight of the ith node.

In the next N-1 lines,each line contains two integers u v,which describes an edge (u,v).

In the next M lines,each line contains three integers u v k,which means an operation asking for the kth minimum weight on the path from node u to node v.

Output

For each operation,print its result.

Example

Input:

 
 

  
  8 5
 
 
8 5
105 2 9 3 8 5 7 7
1 2        
1 3
1 4
3 5
3 6
3 7
4 8
2 5 1
2 5 2
2 5 3
2 5 4
7 8 2 

Output:
2
8
9
105
7 

题意：给出一棵树，询问两点之间第k小的权值是多少

思路：主席树记录从根节点到当前节点出现在1..m（权值按从小到大排序后）中的数（跟poj2104求区间第k小的数一样，这里只不过是书上从根到当前节点），这样当查询时候

在线求出u,v的lca，然后就跟查询区间第k小的值一样了

#include<iostream>
#include<cstdio>
#include<string>
#include<cstring>
#include<vector>
#include<cmath>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<algorithm>
using namespace std;
const int maxn=200010;
const int maxm=maxn*50;
int N,M,m,num;
int w[maxn];
int order[maxn],cnt[maxm],lson[maxm],rson[maxm],T[maxm];
int E[maxn*2],dep[maxn*2],pos[maxn],vis[maxn];
int d[maxn*2][20];
int head[maxn];
struct node
{
    int v,next;
}edge[maxn*2];
void add_edge(int u,int v)
{
    edge[num].v=v;
    edge[num].next=head[u];
    head[u]=num++;
}
void init()
{
    num=0;
    for(int i=0;i<=N;i++)
    {
        vis[i]=0;
        pos[i]=head[i]=-1;
    }
    for(int i=1;i<=N;i++)order[i]=w[i];
    sort(order+1,order+1+N);
    m=unique(order+1,order+1+N)-order-1;
}
//LCA部分
void LCA_dfs(int u,int depth)
{
    E[++num]=u,dep[num]=depth;
    if(pos[u]==-1)pos[u]=num;
    vis[u]=1;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(vis[v])continue;
        LCA_dfs(v,depth+1);
        E[++num]=u;
        dep[num]=depth;
    }
}
void initRMQ(int n)
{
    for(int i=0;i<=n;i++)d[i][0]=i;
    for(int j=1;(1<<j)<=n;j++)
    {
        for(int i=1;i+(1<<j)<=n;i++)
        {
            int x=d[i][j-1],y=d[i+(1<<(j-1))][j-1];
            if(dep[x]<=dep[y])d[i][j]=x;
            else d[i][j]=y;
        }
    }
}
int LCA_query(int u,int v)
{
    int x=pos[u],y=pos[v];
    if(x>y)swap(x,y);
    int k=0;
    while((1<<(k+1))<=y-x+1)k++;
    int l=d[x][k],r=d[y-(1<<k)+1][k];
    if(dep[l]<=dep[r])return E[l];
    else return E[r];
}
//主席树部分
int  build(int l,int r)
{
    int root=num++;
    cnt[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=num++,tmp=newroot;
    int l=1,r=m;
    cnt[newroot]=cnt[root]+val;
    while(l<r)
    {
        int mid=(l+r)>>1;
        if(pos<=mid)
        {
            r=mid;
            lson[newroot]=num++,rson[newroot]=rson[root];
            newroot=lson[newroot],root=lson[root];
        }
        else
        {
            l=mid+1;
            rson[newroot]=num++,lson[newroot]=lson[root];
            newroot=rson[newroot],root=rson[root];
        }
        cnt[newroot]=cnt[root]+val;
    }
    return tmp;
}
int find(int x)
{
    return lower_bound(order+1,order+1+m,x)-order;
}
void dfs_build(int u,int fa)
{
    int pos=find(w[u]);
    T[u]=update(T[fa],pos,1);
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==fa)continue;
        dfs_build(v,u);
    }
}
int query(int left,int right,int lca,int k)
{
    int l=1,r=m;
    int pos=find(w[lca]);
    lca=T[lca];
    while(l<r)
    {
        int mid=(l+r)>>1;
        //注意减掉两倍的lca后，相当于lca这个点没有了，要判断是不是应该加上
        int tmp=cnt[lson[left]]+cnt[lson[right]]-2*cnt[lson[lca]]+(pos>=l&&pos<=mid);
        if(tmp>=k)
        {
            r=mid;
            left=lson[left];
            right=lson[right];
            lca=lson[lca];
        }
        else
        {
            l=mid+1;
            k-=tmp;
            left=rson[left];
            right=rson[right];
            lca=rson[lca];
        }
    }
    return l;
}
int main()
{
    while(scanf("%d%d",&N,&M)!=EOF)
    {
        for(int i=1;i<=N;i++)
            scanf("%d",&w[i]);
        init();
        for(int i=1;i<N;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            add_edge(u,v);
            add_edge(v,u);
        }
        num=0;
        LCA_dfs(1,1);
        initRMQ(num);
        num=0;
        T[0]=build(1,m);
        dfs_build(1,0);
        while(M--)
        {
            int u,v,k;
            scanf("%d%d%d",&u,&v,&k);
            int lca=LCA_query(u,v);
            printf("%d\n",order[query(T[u],T[v],lca,k)]);
        }
    }
    return 0;
}