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<iostream>
#include<cstdio>
#include<vector>
#include<algorithm>
#include<map>
#include<cmath>
#include<set>
#include<cstring>
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<<j)<=n;j++)
{
for(int i=1;i<=n-(1<<j)+1;i++)
{
if(dep[dp[i][j-1]]<dep[dp[i+(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]]<dep[dp[r-(1<<k)+1][k]]) return dp[l][k];
else return dp[r-(1<<k)+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;
}