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 58 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
#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<vector>
#include<map>
#include <bits/stdc++.h>
using namespace std;
const int N = 1e5+100;
typedef long long LL;
int rt[N*20], ls[N*20], rs[N*20], sum[N*20];
int fa[2*N][30], dep[2*N], vis[N];
int a[N], b[N], tot, cnt, head[N], len;
struct node
{
int to, next;
} p[2*N];
void init()
{
memset(head,-1,sizeof(head));
memset(vis,0,sizeof(vis));
cnt=0;
return ;
}
void add(int u,int v)
{
p[cnt].to=v,p[cnt].next=head[u];head[u]=cnt++;
p[cnt].to=u,p[cnt].next=head[v];head[v]=cnt++;
return ;
}
void build(int &o,int l,int r)
{
o= ++tot,sum[o]=0;
if(l==r) return ;
int mid=(l+r)/2;
build(ls[o],l,mid);
build(rs[o],mid+1,r);
return ;
}
void update(int &o,int l,int r,int last,int p)
{
o= ++tot;
ls[o]=ls[last],rs[o]=rs[last];
sum[o]=sum[last]+1;
if(l==r) return ;
int mid=(l+r)/2;
if(p<=mid) update(ls[o],l,mid,ls[last],p);
else update(rs[o],mid+1,r,rs[last],p);
return ;
}
int query(int ss,int tt,int s1,int t1,int l,int r,int cnt)
{
if(l==r) return l;
int tmp=sum[ls[tt]]+sum[ls[ss]]-sum[ls[s1]]-sum[ls[t1]];
int mid=(l+r)/2;
if(tmp>=cnt) return query(ls[ss],ls[tt],ls[s1],ls[t1],l,mid,cnt);
else return query(rs[ss],rs[tt],rs[s1],rs[t1],mid+1,r,cnt-tmp);
}
void dfs(int u,int d,int f,int root)
{
vis[u]=1,dep[u]=d,fa[u][0]=f;
update(rt[u],1,len,root,a[u]);
root=rt[u];
for(int i=head[u];i!=-1;i=p[i].next)
{
int v=p[i].to;
if(vis[v]) continue;
dfs(v,d+1,u,root);
}
return ;
}
void lca(int n)
{
int k=(int)(log(1.0*n)/log(2.0));
for(int i=1;i<=k;i++)
{
for(int j=1;j<=n;j++)
{
fa[j][i]=fa[fa[j][i-1]][i-1];
}
}
return ;
}
int get(int x,int y,int n)
{
if(dep[x]<dep[y]) swap(x,y);
int k=(int)(log(1.0*n)/log(2.0));
int d=dep[x]-dep[y];
for(int i=0;i<=k;i++)
if((d&(1<<i))) x=fa[x][i];
if(x==y) return x;
for(int i=k;i>=0;i--)
{
if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
}
return fa[x][0];
}
int main()
{
int t, n, q;
scanf("%d %d", &n, &q);
for(int i=1; i<=n; i++) scanf("%d", &a[i]), b[i]=a[i];
sort(b+1,b+n+1);
len=unique(b+1,b+n+1)-(b+1);
tot=0;
build(rt[0],1,len);
for(int i=1; i<=n; i++) a[i]=lower_bound(b+1,b+len+1,a[i])-(b);
init();
for(int i=0;i<n-1;i++)
{
int x, y;
scanf("%d %d", &x, &y);
add(x,y);
}
dfs(1,1,0,rt[0]);
lca(n);
while(q--)
{
int l, r, x;
scanf("%d %d %d", &l, &r, &x);
int pos=get(l,r,n);
printf("%d\n",b[query(rt[l],rt[r],rt[pos],rt[fa[pos][0]],1,len,x)]);
}
return 0;
}