SPOJ-COT-Count on a tree

COT-Count on a tree

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小
解题思路：主席树+LCA

#include
#include
#include
#include
#include
#include
#include
#include
using namespace std;
typedef long long LL;
const int MAXN=1e5+5;
int n,m,tot;
int cnt,head[MAXN];
int len,root[MAXN],lson[MAXN*20],rson[MAXN*20],val[MAXN*20];
int num,dep[MAXN*2],ver[MAXN*2],fst[MAXN],dp[MAXN*2][20],fa[MAXN];
bool vis[MAXN];
int a[MAXN],b[MAXN];

struct Edge
{
    int to,nxt;
}e[MAXN*2];

void addedge(int u,int v)
{
    e[cnt].to=v;
    e[cnt].nxt=head[u];
    head[u]=cnt++;
}

void build(int l,int r,int &rt)
{
    rt=++tot;
    val[rt]=0;
    if(l==r) return;
    int mid=(l+r)>>1;
    build(l,mid,lson[rt]);
    build(mid+1,r,rson[rt]);
}

void update(int pre,int &rt,int l,int r,int v)
{
    rt=++tot;
    lson[rt]=lson[pre];rson[rt]=rson[pre];val[rt]=val[pre]+1;
    if(l==r) return;
    int mid=(l+r)>>1;
    if(v<=mid) update(lson[pre],lson[rt],l,mid,v);
    else update(rson[pre],rson[rt],mid+1,r,v);
}

void dfs(int u,int fat,int d)
{
    vis[u]=true;ver[++num]=u;dep[num]=d;fst[u]=num;fa[u]=fat;
    update(root[fat],root[u],1,len,a[u]);
    for(int i=head[u];i!=-1;i=e[i].nxt)
    {
        int to=e[i].to;
        if(!vis[to])
        {
            dfs(to,u,d+1);
            ver[++num]=u;dep[num]=d;
        }
    }
}

void ST(int n)
{
    for(int i=1;i<=n;i++) dp[i][0]=i;
    for(int j=1;(1<for(int i=1;i<=n-(1<1;i++)
        {
            if(dep[dp[i][j-1]]1<<(j-1))][j-1]])
            dp[i][j]=dp[i][j-1];
            else dp[i][j]=dp[i+(1<<(j-1))][j-1];
        }
    }
}

int RMQ(int l,int r)
{
    int k=log(r-l+1)/log(2);
    if(dep[dp[l][k]]1<1][k]]) return dp[l][k];
    else return dp[r-(1<1][k];
}

int LCA(int u,int v)
{
    u=fst[u],v=fst[v];
    if(u>v) swap(u,v);
    int res=RMQ(u,v);
    return ver[res];
}

int query(int ss,int tt,int lca,int lcafa,int l,int r,int k)
{
    if(l==r) return l;
    int mid=(l+r)>>1;
    int tmp=val[lson[ss]]+val[lson[tt]]-val[lson[lca]]-val[lson[lcafa]];
    if(k<=tmp) return query(lson[ss],lson[tt],lson[lca],lson[lcafa],l,mid,k);
    else return query(rson[ss],rson[tt],rson[lca],rson[lcafa],mid+1,r,k-tmp);
}

int main()
{
    while(scanf("%d%d",&n,&m)!=EOF)
    {
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a[i]);
            b[i]=a[i];
        }
        sort(b+1,b+1+n);
        len=unique(b+1,b+1+n)-(b+1);
        for(int i=1;i<=n;i++) a[i]=lower_bound(b+1,b+len+1,a[i])-b;
        cnt=0;
        memset(head,-1,sizeof(head));
        int u,v,k;
        for(int i=1;i<=n-1;i++)
        {
            scanf("%d%d",&u,&v);
            addedge(u,v);addedge(v,u);
        }
        tot=0;
        build(1,len,root[0]);
        num=0;
        memset(vis,false,sizeof(vis));
        dfs(1,0,1);
        ST(2*n-1);
        int lca;
        for(int i=1;i<=m;i++)
        {
            scanf("%d%d%d",&u,&v,&k);
            lca=LCA(u,v);
            printf("%d\n",b[query(root[u],root[v],root[lca],root[fa[lca]],1,len,k)]);
        }
    }
    return 0;
}