백준 1197 - 최소 스패닝 트리 (크루스칼 알고리즘 & 유니온 파인드)

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

1197번: 최소 스패닝 트리

0.

최단 경로 알고리즘에서는 vertex를 기준으로 알고리즘을 수행했다면, 최소 스패닝 트리 알고리즘 (크루스칼, 프림)에서는 edge를 기준으로 알고리즘을 수행한다.  따라서 평소와 다르게 그래프를 이런 식으로 구성하는 편이 낫다. 이런 표현 방식을 edge list라고도 부른다.

 

//생성
vector<tuple<int, int, int>> edge;

//입력 : 비용, 시작, 도착
cin >> u >> v >> w;
edge.emplace_back(w, u, v);

 

edge list가 완성되었으면 크루스칼 알고리즘을 수행한다. 크루스칼 알고리즘은 다음과 같이 수행된다.

 

현재 최소 비용 edge인 min에 대해, 출발 노드와 도착 노드를 각각 u, v라고 하자. 만약 u와 v의 root가 다르다면 cycle이 생기지 않으므로 min을 MST에 추가한다. 이 과정을 MST의 edge 수가 V - 1이 될 때까지 반복한다.

 

이제 문장을 하나씩 해체하고 지식을 습득해보자.

 

1. 현재 최소 비용 edge인 min

우리는 그래프를 edge list의 형태로 입력받았으므로, 이 리스트를 비용 기준으로 오름차순 정렬하면 전처리 O(E log E) 이후 매번 O(1)의 비용으로 최소 비용 간선에 접근할 수 있다. 다행히 c++에서는 sort 함수가 튜플도 잘 정렬해준다. 대신 정렬 최우선 기준이 되는 '비용'을 맨 앞에 저장해야 한다.

 

2. u와 v의 root가 다르다면

유니온 파인드에서 root란 해당 노드가 속한 트리의 root 노드를 의미한다. u와 v의 root를 확인하기 위해서 우리는 find 함수를 사용한다. 유니온 파인드란 union과 find 함수를 통해 여러 트리를 병합하는 기법이다.

 

find는 이런 식으로 구현한다.

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

parent[i]에는 노드 i의 부모 노드가 저장되어있다. 트리의 root 노드는 부모가 존재하지 않으므로 자기자신을(또는 -1을) 부모로 갖는다. 따라서 find 함수에서는 현재 노드가 non-root 노드이면 부모 노드로 이동해 재귀적으로 root를 탐색하며, root에 도달(root[i] 가 i 또는 -1) 하면 그 노드를 반환한다.

 

* path compression 기법을 사용하면 find 함수를 최적화 할 수 있다. 일단 배열에 노드의 부모를 저장하는 대신  root를 저장한다. 구현은 이런 식이다.

int find_improved(int i) {
	if(root[i] == i)
    	return i;
        
    else
    	return root[i] = find_improved(root[i]);
}

사실 이름은 거창하지만 딱 한 줄만 딸깍 수정하면 된다. 재귀적으로 root 노드를 탐색하면서 현재 노드의 root도 즉시 갱신해주면 된다. 예를 들어 1 -> 4 -> 6 -> 7 -> 10으로 이어지는 root가 10인 트리가 있다고 치자. find_improved(1)을 수행하면 트리 구조 대로 1 - 4 - 6 - 7 - 10으로 이동해 마지막에 10을 반환한다. 그러면 10이 재귀적으로 반환되면서 root[7] - root[6] - root[4] - root[1]의 값이 모두 10으로 갱신된다.

 

3. min을 MST에 추가

MST에 추가하는 건 union 함수를 이용해 구현한다. 사실 union 함수도 간단하다. 더 작은 트리를 더 큰 트리에 갖다 붙이면 된다. 여기서 크고 작은 건 각 트리의 root 노드의 level (rank)를 이용해 결정한다. 갖다 붙이는 건 작은 트리 root 노드의 root를 자기 자신이 아닌 큰 트리의 root 노드로 변경하면 된다.

 

만약 두 트리의 크기가 같다면 한 쪽 (일반적으로 왼쪽) 트리에 다른 쪽 트리를 붙이고 level을 1 증가시킨다.

void uni(int x, int y) {
	int xRoot = find(x), yRoot = find(y);

	if (level[xRoot] > level[yRoot])
		root[yRoot] = root[xRoot];
	else if (level[yRoot] > level[xRoot])
		root[xRoot] = root[yRoot];
	else {
		root[yRoot] = root[xRoot];
		level[xRoot]++;
	}
}

MST에 현재 간선 min을 추가했으니, 전체 가중치의 합에 min의 w 값을 더해주자.

 

4. edge 수가 V-1이 될 때까지 반복

edge 수를 기록할 cnt 변수를 사용해 새 edge를 추가할 때마다 cnt를 1씩 증가시킨다.

정점의 수가 V인 트리의 edge 수는 V-1개이므로, cnt == V-1이면 종료한다.

 

전체 코드

int FIND(int i) {
	if (i == root[i])
		return i;
	else
		return root[i] = FIND(root[i]);
}

void UNION(int x, int y) {
	int xRoot = FIND(x), yRoot = FIND(y);

	if (level[xRoot] > level[yRoot])
		root[yRoot] = root[xRoot];
	else if (level[yRoot] > level[xRoot])
		root[xRoot] = root[yRoot];
	else {
		root[yRoot] = root[xRoot];
		level[xRoot]++;
	}
}

void kruskal() {
	for (int i = 1; i <= V; i++)
		root[i] = i;

	int cnt = 0, i = 0, sum = 0;
	while (cnt < V - 1) {
		const auto& nxt = edge[i++];

		auto [w, u, v] = nxt;
		if (FIND(u) == FIND(v))
			continue;

		sum += w;
		UNION(u, v);
		cnt++;
	}

	cout << sum << endl;
}

int main() {
	FASTIO;

	cin >> V >> E;

	int u, v, w;
	while (E) {
		E--;
		cin >> u >> v >> w;
		edge.emplace_back(w, u, v);
	}

	sort(edge.begin(), edge.end());

	kruskal();
}

 

유니온 파인드 기법은 정형화된 코드를 그대로 사용하는 게 편한 듯