[백준] 1761 - 정점들의 거리

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

1761번: 정점들의 거리

 

문제를 보자마자 LCA(log)가 떠올랐다.

그런데 이 문제는 이웃한 정점 간의 거리가 주어지고, 임의의 두 정점 간의 거리를 구하는 문제이다.

 

두 가지 정도 방법이 존재하겠다.

1. 루트 노드부터 모든 정점간의 거리를 dfs로 구하고, dis[a] + dis[b] - 2 * dis[LCA(a, b)]를 구하면 답이다.
2. 두 정점부터 LCA를 찾아가며 거리를 더한다.

2번 먼저 코드로 보겠다.

 

#include <iostream>
#include <vector>
using namespace std;

int depth[40001], dp[40001][20], len[40001][20], n;
vector<vector<pair<int, int>>> v;

int lca(int a, int b) {

	int ret = 0;

	if (depth[a] != depth[b]) {
		if (depth[a] > depth[b]) {
			swap(a, b);
		}

		for (int i = 19; i >= 0; i--) {
			if (depth[a] <= depth[dp[b][i]]) {
				ret += len[b][i];
				b = dp[b][i];
			}
		}
	}

	if (a != b) {
		int pre = 0;
		for (int i = 19; i >= 0; i--) {
            if(dp[a][i] != dp[b][i]) {
				ret += len[a][i] + len[b][i];
				a = dp[a][i];
				b = dp[b][i];
            }
			ret += len[a][i] + len[b][i] - pre;
			pre = len[a][i] + len[b][i];
		}
	}
	return ret;
}

void go(int node, int pre, int h) {
	depth[node] = h;
	dp[node][0] = pre;

	for (auto& i : v[node]) {
		if (i.first != pre) {
			go(i.first, node, h + 1);
			len[i.first][0] = i.second;
		}
	}
}

int main() {

	cin >> n;
	v.resize(n + 1);

	for (int i = 0; i < n - 1; i++) {

		int a, b, c; cin >> a >> b >> c;

		v[a].push_back({ b,c });
		v[b].push_back({ a,c });
	}

	go(1, 0, 1);

	for (int j = 1; j < 20; j++) {
		for (int i = 1; i <= n; i++) {
			dp[i][j] = dp[dp[i][j - 1]][j - 1];
			len[i][j] = len[dp[i][j - 1]][j - 1] + len[i][j - 1];
		}
	}

	int m; cin >> m;

	for (int i = 0; i < m; i++) {
		int a, b; cin >> a >> b;
		cout << lca(a, b) << '\n';
	}

	return 0;
}

 

len 배열만 잘 보면 된다. LCA를 구할때 처럼, len[i][j] = i정점부터 2^j 번째 부모까지 거리이다.

그래서 lca를 구하는 과정에서 깊이를 맞추는 부분에서 더해주었다.

 

그다음은 아래 부분이다.

 

if (a != b) {
	int pre = 0;
	for (int i = 19; i >= 0; i--) {
		if(dp[a][i] != dp[b][i]) {
			ret += len[a][i] + len[b][i];
				a = dp[a][i];
				b = dp[b][i];
        }
		ret += len[a][i] + len[b][i] - pre;
		pre = len[a][i] + len[b][i];
    }
}

 

이 부분 같은 경우에는, 부모가 다를 경우 무조건 그 거리를 지나가야 하므로 더해준다(ret += len[a][i] + len[b][i]).

 

그런데, a와 b의 lca가 첫 번째 부모(i = 0)이고, i가 1일 때를 생각해보자. i가 1일 때도 당연히 부모가 같긴 한데, lca는 아니다. 그래서 이전 부모들끼리(2^0번째 부모와 2^1번째 부모)의 길이를 저장해놓고(pre) 거기서 빼준다.

 

 

 

다음은 1번째 방법이다. 이 부분은 기존 lca를 구하는 코드에서 정점 간의 거리만 구하면 끝이므로 조금 더 간단하다.

 

#include <iostream>
#include <vector>
using namespace std;

int depth[40001], dp[40001][20], dis[40001], n;
vector<vector<pair<int, int>>> v;

int lca(int a, int b) {

	int ret = 0;

	if (depth[a] != depth[b]) {
		if (depth[a] > depth[b]) {
			swap(a, b);
		}

		for (int i = 19; i >= 0; i--) {
			if (depth[a] <= depth[dp[b][i]]) {
				b = dp[b][i];
			}
		}
	}

	ret = a;

	if (a != b) {
		for (int i = 19; i >= 0; i--) {
			if (dp[a][i] != dp[b][i]) {
				a = dp[a][i];
				b = dp[b][i];
			}
			ret = dp[a][i];
		}
	}
	return ret;
}

void go(int node, int pre, int h, int l) {
	depth[node] = h;
	dp[node][0] = pre;
	dis[node] = l;

	for (auto& i : v[node]) {
		if (i.first != pre) {
			go(i.first, node, h + 1, l + i.second);
		}
	}
}

int main() {

	cin >> n;
	v.resize(n + 1);

	for (int i = 0; i < n - 1; i++) {

		int a, b, c; cin >> a >> b >> c;

		v[a].push_back({ b,c });
		v[b].push_back({ a,c });
	}

	go(1, 0, 1, 0);

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

	int m; cin >> m;

	for (int i = 0; i < m; i++) {
		int a, b; cin >> a >> b;
		cout << dis[a] + dis[b] - 2 * dis[lca(a, b)] << '\n';
	}

	return 0;
}