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 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 <bits/stdc++.h>
#define ll long long
const int MAXN=1e5+10;
using namespace std;
ll readll(){
ll x=0,f=1;char ch=getchar();
while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
return x*f;
}
int readint(){
int x=0,f=1;char ch=getchar();
while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();}
while(isdigit(ch))x=x*10+ch-'0',ch=getchar();
return f*x;
}
typedef struct node{
int l,r,ans;
}node;
node d[MAXN*25];
int cnt,rt[MAXN],p[MAXN],n,a[MAXN],q,cnt2;
vector<int>vec[MAXN];
vector<int>vv;
void update(int &x,int y,int l,int r,int vul){
x=++cnt;d[x]=d[y];d[x].ans++;
if(l==r) return ;
int mid=(l+r)>>1;
if(vul<=mid) update(d[x].l,d[y].l,l,mid,vul);
else update(d[x].r,d[y].r,mid+1,r,vul);
}
void querty(int x,int y,int z,int l,int r,int &ans,int k,int lCa){
if(l==r){
ans=l;return ;
}
int mid=(l+r)>>1;
int t=d[d[x].l].ans+d[d[y].l].ans-2*d[d[z].l].ans+(l<=lCa&&lCa<=mid);
if(k<=t) querty(d[x].l,d[y].l,d[z].l,l,mid,ans,k,lCa);
else querty(d[x].r,d[y].r,d[z].r,mid+1,r,ans,k-t,lCa);
}
int b[MAXN<<1],c[MAXN<<1],cnt1,first[MAXN],dp[MAXN<<1][21],pos[MAXN<<1][21],mu[21],mn[MAXN<<1];
void pre_rmq(){
int t=(n<<1);mn[0]=-1;
for(int i=1;i<=t;i++) mn[i]=((i&(i-1))==0)? mn[i-1]+1 : mn[i-1];
for(int i=1;i<=t;i++) dp[i][0]=c[i],pos[i][0]=i;
for(int i=1;i<=20;i++){
for(int j=1;j+mu[i]-1<=t;j++){
if(dp[j][i-1]<=dp[j+mu[i-1]][i-1]) dp[j][i]=dp[j][i-1],pos[j][i]=pos[j][i-1];
else dp[j][i]=dp[j+mu[i-1]][i-1],pos[j][i]=pos[j+mu[i-1]][i-1];
}
}
}
int rmq(int l,int r){
int k=mn[r-l+1];
int t1=dp[l][k];int t2=dp[r-mu[k]+1][k];
if(t1<=t2) return b[pos[l][k]];
else return b[pos[r-mu[k]+1][k]];
}
void dfs(int v,int pre,int dep){
b[++cnt1]=v;c[cnt1]=dep+1;first[v]=cnt1;
cnt++;p[v]=++cnt2;update(rt[cnt2],rt[p[pre]],1,n,a[v]);
for(int i=0;i<vec[v].size();i++){
int u=vec[v][i];
if(u!=pre){
dfs(u,v,dep+1);
b[++cnt1]=v;c[cnt1]=dep+1;
}
}
}
int main(){
mu[0]=1;
for(int i=1;i<=20;i++) mu[i]=(mu[i-1]<<1);
n=readint();q=readint();int u,v,k,u1,v1,Lca,ans;
for(int i=1;i<=n;i++) a[i]=readint(),vv.push_back(a[i]);
for(int i=1;i<n;i++){
u=readint();v=readint();vec[u].push_back(v);
vec[v].push_back(u);
}
sort(vv.begin(),vv.end());
int size=unique(vv.begin(),vv.end())-vv.begin();
for(int i=1;i<=n;i++) a[i]=lower_bound(vv.begin(),vv.begin()+size,a[i])-vv.begin()+1;
dfs(1,0,0);
pre_rmq();
while(q--){
u=readint();v=readint();k=readint();u1=p[u];v1=p[v];Lca=rmq(min(first[u],first[v]),max(first[u],first[v]));
querty(rt[u1],rt[v1],rt[p[Lca]],1,n,ans,k,a[Lca]);
printf("%d
",vv[ans-1]);
}
return 0;
}