백준 14942 개미

https://www.acmicpc.net/problem/14942

LCA를 구하는 과정과 매우 유사하기 때문에, LCA를 모르시는 분이라면 LCA 를 보고 오시는 것을 추천드립니다.

알고리즘은 다음과 같습니다.

  1. LCA를 구하는 탐색(dfs) 과정에서 부모 노드만 저장하는 것이 아니라 부모 노드까지 거리도 저장해 줍니다.

  2. 개미가 자신의 체력을 모두 소진하여 갈 수 있는 노드로 jump 합니다. LCA에서 쿼리를 구하는 방식으로 점프하면 log(n)만에 점프할 수 있습니다.

#include<bits/stdc++.h>
using namespace std;
#define ff first
#define ss second
const int SZ=100005;
const int MXH=18;
typedef pair<int,int> pii;
vector<pii> graph[SZ];
int par[MXH][SZ];
int dst[MXH][SZ];
int level[SZ];
int pw[SZ];
int n;
int dfs(int x, int p, int w, int lv){
	dst[0][x]=w;
	par[0][x]=p;
	level[x]=lv;
	for(auto nxt:graph[x]){
		if(nxt.ff==p) continue;
		dfs(nxt.ff,x,nxt.ss,lv+1);
	}
}
int memo(){
	for(int i=1;i<MXH;i++){
		for(int j=1;j<=n;j++){
			par[i][j]=par[i-1][par[i-1][j]];
			dst[i][j]=dst[i-1][j]+dst[i-1][par[i-1][j]];
		}
	}
}
int lca(int x, int w){
	for(int i=MXH-1;i>=0;i--){
		if(w>=dst[i][x] && par[i][x]!=0){
			w-=dst[i][x];
			x=par[i][x];
		}
	}
	return x;
}
int main(void){
	cin>>n;
	for(int i=1;i<=n;i++) cin>>pw[i];
	for(int i=1;i<n;i++){
		int x,y,z; cin>>x>>y>>z;
		graph[x].push_back({y,z});
		graph[y].push_back({x,z});
	}
	dfs(1,0,0,0);
	memo();
	for(int i=1;i<=n;i++){
		cout<<lca(i,pw[i])<<'\n';
	}
	return 0;
}