백준 1761번 정점들의 거리 C++ 최소 공통 조상, 유니온 파인드

리뷰

 

그래프의 정보가 주어지고 두 노드간의 거리를 구하는 문제, 루트 노드가 1이 고정이 아니여서 직접 구해주어야 한다.

LCA는 개념만 잘 잡는다면 응용 문제에 적용하는 방법이 어렵지 않은 것 같다.

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

 

문제 풀이

1. main 함수에 init, ds, pre_process, query 네개의 함수를 순차적으로 실행시켜 문제를 풀었다.
2. init 함수에서 n값을 받고 n - 1개의 간선 정보를 tree 벡터 배열에 저장한 뒤 Union 작업을 수행해 준다.
3. 1부터 n까지의 노드를 탐색 한 후 자기 자신을 value로 갖는 노드를 찾아 root로 저장한다.
4. root 부터 dfs 함수를 통해 각 노드의 깊이와, 루트 노드로 부터의 거리, 자기 자신의 바로 윗 조상을 찾아준다.
5. pre_process 함수를 통해 모든 노드의 자기 조상에 대한 정보를 dp 형식으로 저장해 준다.
6. query 함수에서 m 값을 입력 받고 총 m개의 쿼리에 대한 출력을 진행해 준다.
7. 두 개의 노드를 입력 받고, lca 함수를 통해 두 노드간 최소 공통 조상을 찾아 변수 par에 해당 노드를 저장한다.
8. 루트부터 a까지의 비용, b까지의 비용을 더한 후 루트부터 par까지의 비용 * 2만큼을 빼준 뒤 출력한다.

참고 사항

자신의 2^j번째 위 조상을 구하는 점화식 : parent[i][j] = parent[parent[i][j - 1]][j - 1]

lca 함수에서 역순으로 탐색할 경우 LOG - 1부터 0까지 탐색해 주어야 한다, LOG부터 탐색해 주어 오답 처리가 한번 됐었다.

이 문제는 입력 값이 적은 편이지만, 웬만하면 ios::sync_with_stdio(0);, cin.tie(0);를 사용해 주는 것이 좋다.

 

정답 코드

#include<iostream>
#include<vector>
#include<algorithm>

using namespace std;

int n, m, root;
const int MAX_N = 40001;
const int LOG = 17;
int parent[MAX_N][LOG];
int depth[MAX_N];
int cost[MAX_N];
int nodes[MAX_N];
vector<pair<int, int>> tree[MAX_N];

int Find(int a) {
	if (nodes[a] == a) return a;
	return nodes[a] = Find(nodes[a]);
}

void Union(int a, int b) {
	int A = Find(a);
	int B = Find(b);
	if (A == B) return;
	nodes[B] = A;
}

void init() {
	cin >> n;
	for (int i = 1; i <= n; i++) nodes[i] = i;
	for (int i = 0; i < n - 1; i++) {
		int a, b, c; cin >> a >> b >> c;
		tree[a].push_back({ b, c });
		tree[b].push_back({ a, c });
		Union(a, b);
	}
	for (int i = 1; i <= n; i++) {
		if (nodes[i] == i) {
			root = i;
			break;
		}
	}
}

void dfs(int node, int p, int d, int c) {
	parent[node][0] = p;
	depth[node] = d;
	cost[node] = c;
	for (pair<int, int> next : tree[node]) {
		int child = next.first, nc = next.second;
		if (child != p) dfs(child, node, d + 1, c + nc);
	}
}

void pre_process() {
	for (int j = 1; j < LOG; j++) {
		for (int i = 1; i <= n; i++) {
			if (parent[i][j - 1] != -1) parent[i][j] = parent[parent[i][j - 1]][j - 1];
		}
	}
}

int lca(int a, int b) {
	if (depth[a] < depth[b]) swap(a, b);
	int diff = depth[a] - depth[b];
	for (int i = 0; i < LOG; i++) {
		if ((diff >> i) & 1) a = parent[a][i];
	}
	if (a == b) return a;
	for (int i = LOG - 1; i >= 0; i--) {
		if (parent[a][i] != parent[b][i]) {
			a = parent[a][i];
			b = parent[b][i];
		}
	}
	return parent[a][0];
}

void query() {
	cin >> m;
	while (m--) {
		int a, b; cin >> a >> b;
		int par = lca(a, b);
		cout << cost[a] + cost[b] - cost[par] * 2 << "\n";
	}
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	init();
	dfs(root, -1, 0, 0);
	pre_process();
	query();
}