[ 백준 / 1240 ] 노드 사이의 거리 (C++)

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

 

난이도 : 골드 5 
알고리즘 유형 : 그래프탐색, BFS, DFS 
풀이 시간 : 15분 + 34분

문제 풀이 

정점과 정점 사이 비용이 주어진 양방향 그래프를 탐색하는 문제이다. 

문제의 풀이 순서는 간단하다.

 

1. 양방향 그래프를 생성한다.

2. `M`만큼 순회하며 정점과 정점 사이 비용을 계산하여 출력한다.

 

위 방법을 사용하여 정점과 정점 사이의 거리(비용)을 구할 수 있다. 하지만 이렇게 입력이 주어질 때마다 정점의 거리를 구하여 출력하게 되면, $정점의 개수 ≤ 1000$이고, $간선의 개수 ≤ 1000$이기 때문에 C++치고 꽤나 긴 시간이 걸린다. 

따라서, 단순히 순회하는 방식이 아닌, LCA 알고리즘을 사용하여 풀게된다면 런타임 시간이 줄어들게 된다.

 

LCA(Lowest Common Ancestor) 알고리즘이란 ? 트리에서 최소 공통 조상을 찾는 알고리즘으로, 이분탐색을 베이스로 찾아나가기 때문에 $O(nlogn)$ 시간복잡도로 문제를 해결할 수 있다.

1. DFS로 깊이와 거리 계산하기
   • 트리의 루트에서 시작하여 DFS를 사용해 각 노드의 깊이와 부모 노드를 저장한다.
   • 또한, 루트에서 각 노드까지의 거리를 dist[] 배열에 저장해둔다. 이렇게 하면 특정 노드까지의 거리를 빠르게 참조할 수 있게 된다.
2. 부모 정보 설정하기 (Sparse Table)
   • 트리에서 LCA를 빠르게 구하기 위해 각 노드의 $2^i$번째 부모 정보를 미리 계산해둔다.
   • 이때, `parent[i][j]`는 노드 $i$의 $2^j$번째 부모를 의미하게 된다. 이는 트리의 특성을 생각해보면 쉽게 이해할 수 있다.
   • 이러한 부모 정보를 설정하면 특정 노드의 조상을 효율적으로 찾을 수 있다.
3. LCA를 이용해 두 노드의 최소 공통 조상 찾기
   • 두 노드의 깊이를 맞춘 후, 동시에 위로 올라가며 최소 공통 조상을 찾는다.
   • 깊이가 다르면 깊은 노드를 더 위로 올려 두 노드의 깊이를 맞추고, 그 후 조상을 비교하며 위로 올라가면서 찾게 된다.
4. 거리 계산
   • 두 노드 u와 v의 거리는 다음과 같이 계산할 수 있다.
     • $거리 = dist[u] + dist[v] * dist[ancestor] - 2$
   • `dist[u]`와 `dist[v]`는 각각 루트에서 u와 v까지의 거리이고, 두 노드의 공통 경로에서 중복되는 거리를 빼주기 위해 LCA의 거리를 두 번 빼게 되면 끝난다.

 

코드 1 : BFS

/*
N개의 노드로 이루어진 트리
M개의 두 노드 쌍을 입력받을 때, 두 노드 사이의 거리 출력

풀이:
bfs, 입력 들어올 때마다 bfs 돌려서 찾기
*/

#include <iostream>
#include <queue>
#include <vector>
#include <cstring>
#define INIT cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);

using namespace std;

int n, m;
int u, v, w;
int visited[1001] = { 0, };


int bfs(int su, int end, vector<vector<pair<int, int>>> graph) {
	queue<pair<int, int>> q;
	visited[su] = 1;
	q.push({ su, 0 });

	while (!q.empty()) {
		u = q.front().first;
		w = q.front().second;
		q.pop();

		for (int i = 0; i < graph[u].size(); i++) {
			v = graph[u][i].first;
			int nw = graph[u][i].second;
			if (!visited[v]) {
				if (v == end)
					return w + nw;
				visited[v] = 1;
				q.push({ v, w + nw });
			}
		}
	}

	return -1;

}

int main() {
	cin >> n >> m;
	vector<vector<pair<int, int>>> graph(n + 1);
	for (int i = 0; i < n - 1; i++) {
		cin >> u >> v >> w;
		graph[u].push_back({ v, w });
		graph[v].push_back({ u, w });
	}

	int start, end;
	while(m--) {
		cin >> start >> end;
		memset(visited, 0, 1001 * sizeof(int));
		cout << bfs(start, end, graph) << '\n';

	}

	return 0;
}

 

코드 2 : LCA

/*
N개의 노드로 이루어진 트리
M개의 두 노드 쌍을 입력받을 때, 두 노드 사이의 거리 출력

풀이:
bfs, 입력 들어올 때마다 bfs 돌려서 찾기

+ LCA 알고리즘 활용하여 풀어보기 : 공통 조상을 미리 구한 뒤, 거리 계산
*/

#include <iostream>
#include <vector>
#include <cstring>
#define INIT cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);

using namespace std;

int n, m;
int u, v, w;
const int MAX = 1001;
const int LOG = 10; // 트리의 높이(log)를 결정할 수 있는 최대 값
int depth[MAX];
int parent[MAX][LOG];
int dist[MAX];
vector<pair<int, int>> graph[MAX];

void dfs(int v, int u, int d) {
	depth[v] = depth[u] + 1;
	parent[v][0] = u;
	dist[v] = d;

	for (auto& edge : graph[v]) {
		int next = edge.first;
		int weight = edge.second;
		if (next != u) {
			dfs(next, v, d + weight);
		}
	}
}

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

int lca(int u, int v) {
	if (depth[u] < depth[v]) swap(u, v);

	for (int i = LOG - 1; i >= 0; i--) {
		if (depth[u] - (1 << i) >= depth[v]) {
			u = parent[u][i];
		}
	}

	if (u == v) return u;

	for (int i = LOG - 1; i >= 0; i--) {
		if (parent[u][i] != parent[v][i]) {
			u = parent[u][i];
			v = parent[v][i];
		}
	}

	return parent[u][0];
}

int main() {
	INIT;
	cin >> n >> m;

	for (int i = 0; i < n - 1; i++) {
		cin >> u >> v >> w;
		graph[u].push_back({ v, w });
		graph[v].push_back({ u, w });
	}

	depth[0] = -1;
	dfs(1, 0, 0); // 루트 정점를 1로 가정함
	setParent(n);

	while (m--) {
		int u, v;
		cin >> u >> v;
		int ancestor = lca(u, v);
		cout << dist[u] + dist[v] - 2 * dist[ancestor] << '\n';
	}

	return 0;
}

배운 점 

LCA 알고리즘을 적용한 지 오래되어 다시 복기하느라 찾아보았다. 코딩테스트에 많이 사용되는 알고리즘은 아니지만, 최적화가 필요한 경우에는 알아두면 매우 좋은 알고리즘인 것 같다.

이 기회에 확실히 알아가도록 다시 공부해야겠당