BOJ 13511 트리와 쿼리 2

1 개요[ | ]

BOJ 13511 트리와 쿼리 2
2 같이 보기[ | ]

3 C++[ | ]

#include <bits/stdc++.h>
using namespace std;

int N, M;
int d[100001];
int p[100001][18];
long long len[100001];
vector<pair<int,int>> v[100001];

void init(int cur) {
	for (const auto& i: v[cur]) {
		if (d[i.second] == -1) {
			d[i.second] = d[cur] + 1;
			len[i.second] = len[cur] + i.first;
			p[i.second][0] = cur;
			init(i.second);
		}
	}
}

int lca(int a, int b) {
	if (d[a] < d[b]) swap(a, b);
	int diff = d[a] - d[b];
	int j = 0;
	while (diff) {
		if (diff % 2) a = p[a][j];
		diff /= 2;
		j++;
	}
	if (a == b) return a;
	for (int j=17; j>=0; j--) {
		if (p[a][j] != -1 && p[a][j] != p[b][j]) {
			a = p[a][j];
			b = p[b][j];
		}
	}
	a = p[a][0];
	return a;
}
int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> N;
	for (int i=0; i<N-1; i++) {
		int a, b, cost;
		cin >> a >> b >> cost;
		v[a].push_back({ cost, b });
		v[b].push_back({ cost, a });
	}
	memset(p, -1, sizeof(p));
	memset(d, -1, sizeof(d));
    d[1] = 0;
	init(1);
	for (int j=0; j<17; j++) {
		for (int i = 2; i <= N; i++) {
			if (p[i][j] != -1) {
				p[i][j + 1] = p[p[i][j]][j];
			}
		}
	}
	cin >> M;
	while (M--) {
		int ch, a, b, k;
		cin >> ch >> a >> b;
		int pnode = lca(a, b);
		if (ch == 1) {
			cout << len[a]+len[b] - 2*len[pnode] << '\n';
			continue;
		}
		cin >> k;
		k--;
		if (d[a] - d[pnode] >= k) {
			for (int i = 0; k > 0; i++) {
				if (k % 2) a = p[a][i];
				k /= 2;
			}
			cout << a << "\n";
			continue;
		}
		k -= d[a] - d[pnode];
		k = d[b] - d[pnode] - k;
		for (int i=0; k>0; i++) {
			if (k%2 != 0) b = p[b][i];
			k /= 2;
		}
		cout << b << '\n';
	}
}