[백준 1197] 최소 스패닝 트리 (C++)

이 문제는 주어진대로, 최소 비용 생성 트리(MST)의 비용을 구하는 문제이다.

MST란, 무방향 그래프의 부분 그래프 중 모든 정점을 포함하는 트리이다. 즉, 부분 그래프 중에서 사이클이 없는 무방향 연결 그래프가 MST란 것이다.

MST를 만드는 알고리즘은 크루스칼과 프림 알고리즘이 있다.
먼저, 크루스칼 알고리즘은 간선을 비용 순으로 정렬하고 비용이 가장 적은 간선부터 선택해나가는 알고리즘이다. 이때, 사이클이 생긴다면 선택하지 않도록 해야하고, 정점 갯수가 $V$개 일 때, 간선을 $V - 1$개 선택하면 알고리즘을 종료해주면 된다.

사이클이 발생하는지 여부는 Flood-Fill이나 Union-Find 알고리즘으로 확인이 가능하다.
Flood-Fill은 BFS나 DFS로 모든 정점을 방문해가며 색을 칠해나가는 알고리즘인데, 색이 칠해진 정점이 방문했다는 것은 재방문을 했다는 의미이므로 사이클이 있음을 확인할 수 있다. 이 알고리즘의 시간 복잡도는 $O(V)$인데, 간선마다 해당 간선을 선택했을 때, 사이클이 생기는지 여부를 판단해야 하므로 시간 복잡도는 $O(VE)$가 된다. 따라서, 시간 초과가 발생하므로 Flood-Fill은 사용하지 못한다.

다음으로, Union-Find 알고리즘은 정점들을 어떤 대푯값을 기준으로 그룹화하는 알고리즘이다. 따라서, 같은 집합(그룹)이라면 같은 대표 원소를 갖도록 관리하는 것이 핵심이다. 따라서, 대표 원소를 저장할 배열 parent를 미리 준비해 두는 과정이 선행되어야 한다. 처음에는 각 원소들의 대표 원소는 자기 자신이 될 것이므로 parent[i] = i가 되도록 초기화를 진행해 두어야 한다. 그럼, Union-Find 알고리즘의 연산에 대해 살펴보자.

이 알고리즘에서 연산은 Union과 Find로 총 2가지의 연산이 존재한다.

Union연산은 원소 x, y가 주어질 때, 두 원소가 속한 집합을 합치는 연산이다. 따라서, Find연산으로 대표 원소를 찾아, 두 원소의 대표 원소가 같지 않다면 둘 중 하나를 다른 집합에 속하도록 해 대표 원소를 변경해 주어야한다. 이를 코드로 구현하면 다음과 같다.

void merge(int x, int y) {
	x = find(x); // x의 대표 원소 대입
	y = find(y); // y의 대표 원소 대입
	if (x == y) return;
	parent[y] = x;
}

하지만, 이렇게 구현하면 parent 배열은 항상 대표 원소를 저장한다고 보장할 수 없어진다. 왜냐하면, 마지막의 parent[y] = x를 통해 대표 원소 1개의 parent배열 값을 변경했기 때문이다. 좀 더 자세히 설명하자면, parent[y]를 대표 원소로 가졌던 원소들은 기존의 parent[y]를 parent 배열 값으로 가지는데, parent[y]값만이 x로 변경되었기 때문이다.

이러한 불일치로 인해, parent배열이 항상 대표 원소를 기억한다는 것을 보장할 수 없다. 오직, parent[x]와 x가 동일할 때만 대표 원소를 기억하는 것을 보장할 수 있게 된다.

Find는 주어진 원소 x의 대표 원소를 찾는 연산으로, x와 parent[x]가 동일하다면 x가 대표 원소이므로 x를 반환한다. x와 parent[x]가 다르다면, parent[x]가 x의 대표 원소임을 보장할 수 없기 때문에, x와 parent[x]가 같아질 때 까지 재귀적으로 재탐색을 해나가야 한다. 이를 코드로 구현하면 다음과 같다.

int find(int x) {
	if (parent[x] == x) return x;
	return parent[x] = find(parent[x]);
}

크루스칼 알고리즘에서 방문 가능한 정점들을 하나의 집합으로 묶어 관리하는 방식으로 Union-Find 알고리즘을 응용하여 사이클 여부를 확인할 수 있다. 이는 거의 상수시간에 사이클을 확인할 수 있게 되므로 간선마다 사이클 여부를 판단하므로 시간복잡도가 $O(E)$가 된다. 그리고, 간선들을 정렬해야하므로 $O(ElogE)$의 시간 복잡도가 필요하여, 전체 시간복잡도는 $O(ElogE)$가 되어 시간 내에 정답을 구할 수 있다.

크루스칼 알고리즘을 이용하여 문제를 해결해보자.

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

int v, e, parent[10'001];
tuple<int, int, int> edge[100'001]; // {비용, 정점1, 정점2}

int find(int x) {
	if (parent[x] == x) return x;
	return parent[x] = find(parent[x]);
}

void merge(int x, int y) {
	x = find(x);
	y = find(y);
	if (x == y) return;
	parent[y] = x;
}

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

	cin >> v >> e;
	for (int i = 0, a, b, c; i < e; i++) {
		cin >> a >> b >> c;
		edge[i] = {c, a, b};
	}

	sort(edge, edge + e);
	for (int i = 1; i <= v; i++) parent[i] = i;

	int cnt = 0;
	long long cost = 0;
	for (int i = 0; i < e; i++) {
		int a, b, c;
		tie(c, a, b) = edge[i];
		if (find(a) == find(b)) continue;
		merge(a, b);
		cost += c;
		cnt++;
		if (cnt == v - 1) break;
	}
	cout << cost;
}

다음으로 프림 알고리즘을 살펴보자. 이 알고리즘은 임의의 정점 $V$에서 최소 신장 트리에 속하지 않은 정점을 연결하는 간선들 중 비용이 최소인 것을 최소 신장 트리에 추가해나간다. 따라서, 우선순위 큐를 이용해 구현이 가능하다.

먼저, 정점 $V$에서 인접 정점들로의 간선을 곧바로 구해야하므로 인접 리스트로 그래프를 표현해야 한다. 그리고, 정점이 최소 신장 트리에 속한 정점인지 확인하기 위한 배열 chk를 도입해 $O(1)$로 판단하자.

인접 리스트에는 {비용, 인접 정점}을 저장해야하며, 우선순위 큐에는 {비용, 정점1, 정점2}를 저장해야 한다.

시작 정점은 임의로 잡아도 되기 때문에 정점 $1$로 잡자. 그리고 정점 $1$을 최소 비용 생성트리에 추가했다는 의미로 chk[1] = 1로 설정한다. 그리고 정점 $1$에서 인접 정점으로의 간선을 우선순위 큐에 {비용, 1, 인접 정점} 형태로 추가한다.

이제, 우선순위 큐에서 최소 비용인 간선을 꺼내고, MST에 속하는 정점들을 연결하는지 확인한다. 먼저, 우선순위 큐의 원소 정보 {비용, 정점1, 정점2}를 각각 c, a, b에 옮겨 담는다고 가정하자. 이때, 우리는 원하는 간선을 선택하면 추가된 새로운 정점에서 인접 정점으로의 간선을 우선순위 큐에 추가할 것이다. 그렇다면, 우선순위 큐에 원소를 삽입하는 형태를 고려하면, 항상 a는 기존에 MST에 속했던 정점이고 b는 MST에 새롭게 추가된 정점임을 알 수 있다.

따라서, 우선순위 큐에서 원소를 꺼냈을 때, 정점 b가 MST에 속한다면 그대로 해당 간선을 우선순위 큐에서 제거하면 되고, 그렇지 않다면 정점 b가 MST에 속하게 되었다는 의미로 chk[b] = 1로 설정하고 해당 간선을 추가한다. 이 문제에서는 간선을 다 기억할 필요 없이 비용 ans와 MST에 추가된 간선 갯수 cnt만 증가시키면 된다.

이제, 정점b에서 인접 정점으로의 간선을 우선순위에 추가해주고, 간선을 $V - 1$개 선택할 때 까지 이 과정을 반복해주면 된다.

이를 코드로 구현하면 다음과 같다.

#include <bits/stdc++.h>
#define X first
#define Y second
using namespace std;

int v, e;
vector<pair<int, int>> adj[10001]; // {비용, 정점}
priority_queue<tuple<int, int, int>,
			   vector<tuple<int, int, int>>,
			   greater<>> pq; // {비용, 정점1, 정점2}
bool chk[10001];

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

	cin >> v >> e;
	for (int i = 0, a, b, c; i < e; i++) {
		cin >> a >> b >> c;
		adj[a].push_back({c, b});
		adj[b].push_back({c, a});
	}

	int cnt = 0;
	long long ans = 0;
	chk[1] = 1;
	for (auto nxt : adj[1]) pq.push({nxt.X, 1, nxt.Y});

	while (cnt < v - 1) {
		int c, a, b;
		tie(c, a, b) = pq.top();
		pq.pop();
		if (chk[b]) continue;
		chk[b] = 1;
		ans += c;
		cnt++;

		for (auto nxt : adj[b]) {
			if (!chk[nxt.Y]) pq.push({nxt.X, b, nxt.Y});
		}
	}
	cout << ans;
}

크루스칼 알고리즘이 구현이 편리하므로 MST문제는 크루스칼로 구현하는 것이 좋아보인다.