백준13511 트리와 쿼리2

작성일 | In

문제 링크

  • http://icpc.me/13511

사용 알고리즘

  • LCA

풀이

Sparse Table을 이용해 LCA를 구해준 뒤, 비슷한 방법으로 경로를 타고 가면 k번째 정점을 알 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<ll, ll> p;

vector<p> g[101010];
ll dp[101010][22];
ll dist[101010];
ll dep[101010];
ll n;
bool chk[101010];

void dfs(ll now, ll d, ll c){
	chk[now] = 1;
	dep[now] = d;
	dist[now] = c;
	for(auto i : g[now]){
		ll nxt = i.first;
		if(chk[nxt]) continue;
		ll nxtCost = i.second + c;
		dp[nxt][0] = now;
		dfs(nxt, d+1, nxtCost);
	}
}

void make(){
	for(int j=1; j<22; j++){
		for(int i=1; i<=n; i++){
			dp[i][j] = dp[dp[i][j-1]][j-1];
		}
	}
}

ll getLca(ll u, ll v){
	if(dep[u] > dep[v]) swap(u, v);
	ll diff = dep[v] - dep[u];
	for(int i=0; diff; i++){
		if(diff & 1) v = dp[v][i];
		diff >>= 1;
	}
	if(u == v) return u;
	for(int i=21; i>=0; i--){
		if(dp[u][i] != dp[v][i]) u = dp[u][i], v = dp[v][i];
	}
	return dp[u][0];
}

ll getDist(ll u, ll v){
	ll lca = getLca(u, v);
	return dist[u] + dist[v] - 2*dist[lca];
}

ll getkth(ll u, ll v, ll k){
	ll lca = getLca(u, v);
	ll diff = dep[u] - dep[lca];
	k--;
	if(diff >= k){
		for(int i=0; k; i++){
			if(k & 1) u = dp[u][i];
			k >>= 1;
		}
		return u;
	}

	k = dep[v] - dep[lca] + diff - k;
	for(int i=0; k; i++){
		if(k & 1) v = dp[v][i];
		k >>= 1;
	}
	return v;
}

int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	cin >> n;
	for(int i=1; i<n; i++){
		ll s, e, x; cin >> s >> e >> x;
		g[s].push_back({e, x});
		g[e].push_back({s, x});
	}
	dfs(1, 1, 0); make();

	int m; cin >> m;
	while(m--){
		int op; cin >> op;
		if(op == 1){
			ll a, b; cin >> a >> b;
			cout << getDist(a, b) << "\n";
		}else{
			ll a, b, c; cin >> a >> b >> c;
			cout << getkth(a, b, c) << "\n";

		}
	}
}