[알고리즘] MST(크루스칼과 프림 알고리즘)

이번 포스팅에서는 최소 신장 트리(MST) 알고리즘을 알아보겠습니다. 최소 신장 트리는 난이도가 높은 알고리즘에 속합니다. 따라서 그래프 이론, 트리의 개념, 서로소 집합(Union-Find), 정렬, 우선순위 큐에 대한 이해가 선행되어야 합니다. 따라서 기본적인 알고리즘을 먼저 이해한 후에 공부하는 것을 권장합니다.

최소 신장 트리 알고리즘은 크게 크루스칼 알고리즘과 프림 알고리즘이 있습니다. 이 글에서는 크루스칼(Kruskal's algorithm), 프림 알고리즘을 소개하고 활용에 대해서 알아보겠습니다.

 

1. 최소 신장 트리란?

V개의 정점과 E개의 간선이 양방향 그래프 형태로 연결되어 있다고 가정해봅시다. 또한 이 그래프는 연결그래프입니다. (연결그래프란 소외된 정점이 없다는 것을 의미합니다. 즉, 모든 정점이 연결되어 있습니다.)

신장트리란?

먼저 신장 트리가 무엇인지 알아야 합니다. 신장트리란 하위 그래프 중에서 모든 정점을 포함하는 트리를 말합니다. 

 

최소 신장 트리란?

위의 신장트리 중에서 간선의 가중치합이 최소인 신장 트리입니다.

이처럼 모든 정점을 최소 비용을 들여 연결해야 하는 상황에서 최소 신장트리를 이용할 수 있습니다.

 

2. 크루스칼 알고리즘

크루스칼 알고리즘은 간선을 미리 크기순으로 정렬해놓고 비교하며 V-1개의 간선이 연결될 때까지 간선을 추가하는 알고리즘입니다. 알고리즘의 순서는 다음과 같습니다.

 

크루스칼 알고리즘 순서

1. 간선들을 가중치를 기준으로 오름차순 정렬한다.
2. 정렬된 간선 들을 탐색하면서 두 정점이 다른 집합일 경우에만 신장 트리에 추가합니다. 
3. 추가된 간선이 V-1개라면 종료합니다.

 

알고리즘 시뮬레이션

 

 

 

알고리즘 구현

이를 코드로 구현하면 간단하게 아래처럼 구현할 수 있습니다. 

tuple<int, int, int> edge[10001];
int v, e;

for (int i = 0; i < e; i++) {
	int a, b, cost;
	tie(cost, a, b) = edge[i];
	if (flood_fill(a, b)) continue;
	cout << cost << " " << a << " " << b << "\n";
	ans += cost;
	cnt++;
	if (cnt == v - 1) break;
}

위의 코드에서 두 정점이 같은 집합인지 확인하는 함수를 flood_fill이라고 했습니다. 다음으로 이 함수를 정의해야 합니다. 이 방법은 크게 두 가지가 있습니다.

 

[원시적인 크루스칼]

먼저 원시적인 완전 탐색 방식입니다. BFS를 통한 Flood Fill을 수행하면 정점 a에서 b로 갈 때 최초 b 방문시 최소 거리임이 보장되므로 O(V)의 시간복잡도로 탐색할 수 있습니다. 코드는 아래와 같습니다.

bool flood_fill(int start, int end) {
	bool visited[10001] = { false, };
	queue<int> q;

	visited[start] = true;
	q.push(start);
	while (!q.empty()) {
		int cur = q.front();
		q.pop();
		for (int& nxt : mst_edge[cur]) {
			if (nxt == end) return true;
			if (visited[nxt]) continue;
			visited[nxt] = true;
			q.push(nxt);
		}
	}
	return false;
}

Flood Fill을 이용할 경우 몇 가지 주의사항이 있습니다.

먼저 신장 트리에 포함된 간선들을 따로 관리해야 합니다. 아래의 전체 코드에서 알 수 있듯이 mst_edge라는 배열을 만들어 현재까지 신장 트리에 포함된 간선들을 따로 저장했습니다.

다음으로 Flood Fill함수는 현재까지 신장 트리에 포함된 간선들만으로 두 정점 a에서 b로 갈 수 있는지 탐색하는 함수입니다. 이 함수는 BFS와 똑같습니다. 방문체크를 까먹으면 무한루프에 걸리므로 이것만 신경써주면 됩니다.

한 번의 탐색마다 O(V)의 시간이 걸리고 매 간선마다 수행해야 하므로 BFS를 이용한 Flood Fill방식은 O(VE)가 됩니다. 따라서 정렬까지 고려한 전체 시간복잡도는 O(ElogE + VE)입니다. 

Flood Fill의 전체 코드는 아래와 같습니다.

#include<iostream>
#include<tuple>
#include<algorithm>
#include<queue>
#include<vector>
using namespace std;


int cnt, ans;
//정점, 간선 개수
int v, e;
//cost, node1, node2
tuple<int, int, int> edge[100001];
//MST에 편입된 간선들을 관리
vector<int> mst_edge[10001];

bool flood_fill(int start, int end) {
	bool visited[10001] = { false, };
	queue<int> q;

	visited[start] = true;
	q.push(start);
	while (!q.empty()) {
		int cur = q.front();
		q.pop();
		for (int& nxt : mst_edge[cur]) {
			if (nxt == end) return true;
			if (visited[nxt]) continue;
			visited[nxt] = true;
			q.push(nxt);
		}
	}
	return false;
}

int main(void) {
	cin >> v >> e;
	//input
	for (int i = 0; i < e; i++) {
		int a, b, cost;
		cin >> a >> b >> cost;
		edge[i] = { cost, a, b };
	}
	sort(edge, edge + e);
	
	cout << "mst_edge: \n";
	for (int i = 0; i < e; i++) {
		int a, b, cost;
		tie(cost, a, b) = edge[i];
		if (flood_fill(a, b)) continue;
		cout << cost << " " << a << " " << b << "\n";
		mst_edge[a].push_back(b);
		mst_edge[b].push_back(a);
		ans += cost;
		cnt++;
		if (cnt == v - 1) break;
	}
	cout << ans;
	return 0;
}

하지만 E 최대 V-1개이므로 정렬에서 O(ElogE)의 시간복잡도를 가진다고 하더라도 결국 O(VE)는 O(V²)에 수렴하게 됩니다. 따라서 이를 개선한 방법이 필요합니다.

 

[유니온 파인드를 이용한 크루스칼]

유니온 파인드를 이용하면 따로 MST에 편입된 간선들을 저장하지 않고 분리 집합 여부를 O(1)에 알 수 있습니다. is_dif함수는 두 노드가 같은 집합인지 확인하고 다른 집합이라면 하나의 집합에 포함시킵니다. 코드로 구현하면 아래와 같습니다.

vector<int> par(10000, -1);

int Find(int x) {
	return par[x] < 0 ? x : par[x] = Find(par[x]);
}

bool is_dif(int x, int y) {
	x = Find(x);
	y = Find(y);
	if (x == y) return 0;
	if (par[x] == par[y]) par[x]--;
	if (par[x] < par[y]) par[x] = y;
	else par[y] = x;
	return 1;
}

위의 코드는 weighted union find, 경로 최적화, 랭크 최적화를 모두 적용하여 유니온 파인드를 구현한 코드입니다. 집합의 관한 모든 연산은 O(1)이므로 유니온 파인드를 이용한 MST 알고리즘은 O(ElogE)입니다.

유니온 파인드를 이용하여 구현한 MST의 전체 코드는 아래와 같습니다.

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

using namespace std;

tuple<int, int, int> edge[100001];
vector<int> par(10000, -1);
int N, e, v;

int Find(int x) {
	return par[x] < 0 ? x : par[x] = Find(par[x]);
}

bool is_dif(int x, int y) {
	x = Find(x);
	y = Find(y);
	if (x == y) return 0;
	if (par[x] == par[y]) par[x]--;
	if (par[x] < par[y]) par[x] = y;
	else par[y] = x;
	return 1;
}

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

	for (int i = 1; i <= N; i++) par[i] = i;

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

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

 

3. 프림 알고리즘

프림 알고리즘은 우선순위 큐를 사용합니다. 

프림 알고리즘 순서

 

시뮬레이션

 

알고리즘 구현

int MST(int start) {
	
	visit[start] = true;
	for (auto& nxt : adj[start]) pq.push({ nxt.F, start, nxt.S });

	while (cnt < V - 1) {
		int cost, a, b;
		tie(cost, a, b) = pq.top(), pq.pop();
		if (visit[b]) continue;
		cout << "MST정점: (" << a << " " << b << "), 비용: " << cost << "\n";
		visit[b] = true;
		ans += cost;
		cnt++;
		for (auto& nxt : adj[b]) if(!visit[nxt.second]) pq.push({ nxt.F, b, nxt.S });
	}
	return ans;
}

 

전체 코드

#include<iostream>
#include<queue>
#include<vector>
#include<tuple>
#define F first
#define S second

using namespace std;

int V, E;
int ans, cnt;
priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<tuple<int, int, int>>> pq;

bool visit[100001];
vector<pair<int, int>> adj[100001];

int MST(int start) {
	
	visit[start] = true;
	for (auto& nxt : adj[start]) pq.push({ nxt.F, start, nxt.S });

	while (cnt < V - 1) {
		int cost, a, b;
		tie(cost, a, b) = pq.top(), pq.pop();
		if (visit[b]) continue;
		cout << "MST정점: (" << a << " " << b << "), 비용: " << cost << "\n";
		visit[b] = true;
		ans += cost;
		cnt++;
		for (auto& nxt : adj[b]) if(!visit[nxt.second]) pq.push({ nxt.F, b, nxt.S });
	}
	return ans;
}

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

	cin >> V >> E;
	for (int i = 0; i < E; i++) {
		int u, v, c;
		cin >> u >> v >> c;
		adj[u].push_back({ c, v });
		adj[v].push_back({ c, u });
	}
	ans = MST(1);
	cout << ans;
	return 0;
}

 

일반적으로 유니온 파인드만 다룰줄 알면 크루스칼이 훨씬 간단합니다. 하지만 STL에 능숙하다면 우선순위 큐를 이용한 프림 알고리즘으로 간결하게 구현할 수 있습니다. 결론적으로는 두 알고리즘을 모두 익히고 상황에 따라 자유자제로 사용하는 것을 권장합니다. 구현을 마치면 [백준 11383] 최소 신장 트리 문제로 확인해 보세요. 간단한 MST를 구현하는 알고리즘입니다.

 

4. 활용

앞의 최소 신장 트리 소개에서 말했듯이, 이 알고리즘은 모든 정점을 최소 비용으로 연결해야 하는 상황에서 유용하게 쓰입니다. 실생활에서 쓰이는 경우를 몇 가지만 소개하면 다음과 같습니다.

도시 분할 계획

모든 도시를 연결하면서 도로의 길이가 최소가 되도록 할 때 MST를 쓸 수 있습니다. 대표적인 문제로는 [백준 1647] 도시 분할 계획이 있습니다.

이와 비슷하게 전기 회로, 통신, 배관, 네트워크 구축에서도 모든 정점들을 연결하면서 비용이 최소가 되도록 할 때 MST를 사용할 수 있습니다.

항상 모든 알고리즘 문제가 그렇듯, 문제를 읽으며 이 알고리즘을 사용해야만 하는 정당성을 생각해야 합니다. MST알고리즘도 MST를 써야만 하는 정당성을 떠올리는 것을 목표로 여러 문제를 풀어봐야 합니다.