Seletiva IOI 2017

arvore.cpp

// Ivan Carvalho
// Árvore- Seletiva IOI - OBI 2017
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 5*1e5 + 10;
const int MAXL = 20;
vector<int> grafo[MAXN],esq,dir,val;
int ancestral[MAXN][MAXL],raiz[MAXN],n,q,segPtr,valor[MAXN],nivel[MAXN];
void utilidade(){
	esq.push_back(0);
	dir.push_back(0);
	val.push_back(0);
}
void update(int novo,int velho,int left,int right,int alvo){
	val[novo] = val[velho] + 1;
	esq[novo] = esq[velho];
	dir[novo] = dir[velho];
	if(left == right) return;
	int mid = (left+right)/2;
	if(alvo <= mid){
		esq[novo] = ++segPtr;
		utilidade();
		update(esq[novo],esq[velho],left,mid,alvo);
	}
	else{
		dir[novo] = ++segPtr;
		utilidade();
		update(dir[novo],dir[velho],mid+1,right,alvo);
	}
}
void dfs(int v,int p){
	nivel[v] = nivel[p] + 1;
	raiz[v] = ++segPtr;
	utilidade();
	update(raiz[v],raiz[p],1,MAXN,valor[v]);
	ancestral[v][0] = p;

	for(int i = 1;i<MAXL;i++){
		int u = ancestral[v][i-1];
		if(u != 0){
			ancestral[v][i] = ancestral[u][i-1];
		}
	}
	for(int u : grafo[v]){
		if(u != p) dfs(u,v);
	}
}
int LCA(int u,int v){
	if(nivel[u] < nivel[v]) swap(u,v);
	for(int i = MAXL - 1;i>=0;i--){
		int prox = ancestral[u][i];
		if(prox != 0 && nivel[prox] >= nivel[v]){
			u = prox;
		}
	}
	if(u == v) return u;
	for(int i = MAXL - 1;i>=0;i--){
		if(ancestral[u][i] != 0 && ancestral[v][i] != 0 && ancestral[u][i] != ancestral[v][i]){
			u = ancestral[u][i];
			v = ancestral[v][i];
		}
	}
	return ancestral[u][0];
}
int kth(int r1,int r2,int r3,int r4,int left,int right,int resta){
	if(left == right) return left;
	int mid = (left+right)/2;
	int davez = val[esq[r1]] + val[esq[r2]] - val[esq[r3]] - val[esq[r4]];
	if(resta <= davez) return kth(esq[r1],esq[r2],esq[r3],esq[r4],left,mid,resta);
	return kth(dir[r1],dir[r2],dir[r3],dir[r4],mid+1,right,resta - davez);
}
int main(){
	scanf("%d %d",&n,&q);
	utilidade();
	for(int i = 1;i<=n;i++){
		scanf("%d",&valor[i]);
	}
	for(int i = 1;i<n;i++){
		int u,v;
		scanf("%d %d",&u,&v);
		grafo[v].push_back(u);
		grafo[u].push_back(v);
	}
	dfs(1,0);
	for(int vez = 1;vez<=q;vez++){
		int k,u,v,pai;
		scanf("%d %d %d",&k,&u,&v);
		pai = LCA(u,v);
		printf("%d\n",kth(raiz[u],raiz[v],raiz[pai],raiz[ancestral[pai][0]],1,MAXN,k));
	}
	return 0;
}