프로그래머스 - 섬 연결하기

두 Vertex를 잇는 Edge의 cost가 주어졌고, 해당 cost가 최소인 MST를 만드는 것이 문제의 목표이다.

 

MST를 구하는 방법은 여러 개가 있는데, 이 중 Kruskal-algorithm을 사용하여 풀어보도록 하겠다. 

 

크루스칼 알고리즘은 union-find 알고리즘을 이용하여 MST를 구하는 알고리즘으로 잘 알려져 있다. 

 

코드

#include <string>
#include <vector>
#include <algorithm>
 
using namespace std;
 
int getRoot(int idx, vector<int> parent)
{
    if (parent[idx] == idx) return (idx);
    else return getRoot(parent[idx], parent);
}
 
bool find(vector<int> parent, int a, int b)
{
    a = getRoot(a, parent);
    b = getRoot(b, parent);
    if (a == b) return true;
    return false;
}
 
void my_union(vector<int> &parent, int a, int b)
{
    a = getRoot(a, parent);
    b = getRoot(b, parent);
    if (a < b) parent[b] = a;
    else parent[a] = b;
}
 
bool comp(vector<int> a, vector<int> b)
{
    return a[2] < b[2];
}
 
int solution(int n, vector<vector<int>> costs) {
    int answer = 0;
    sort(costs.begin(), costs.end(), comp);
    vector<int> parent(n);
    for(int i=0; i<n; i++)
        parent[i] = i;
    for(int i=0; i<costs.size(); i++)
    {
        if (!find(parent, costs[i][0], costs[i][1]))
        {
            my_union(parent, costs[i][0], costs[i][1]);
            answer += costs[i][2];
        }
    }
    return answer;
}

크루스칼 알고리즘은 cost가 최소인 Edge들 부터 연결하는 것이기 때문에 우선 costs 배열을 정렬시켜 놓는다. 

 

union-find 알고리즘을 구현하기 위해

1. 부모 노드를 저장할 vector<int> parent를 Vertex 개수만큼 만들어 놓았다. 

2. parent[i] = i로 각각의 vertex의 부모 노드를 초기화 시켜 놓는다. 

3. find 함수는 두 개의 vertex 번호를 받아 각 vertex의 부모가 누군지 확인하여 부모가 서로 다르다면 false를 반환하는 함수이다. (부모가 서로 다르다는 의미는 현재 서로 연결이 되어 있지 않다는 의미이다.)

4. my_union 함수는 두 개의 vertex 번호를 받아 각 vertex의 부모를 구하여 부모의 vertex 번호가 더 낮은 것이 다른 vertex의 부모가 되도록 한다. 

5. getRoot 함수는 Vertex 번호를 받아 부모를 찾아주는 함수로, 만약 Vertex 번호가 부모의 번호와 같다면(즉, 자신이 부모노드라면) 자신의 번호를 반환하고, 아니면 getRoot를 다시 호출하여 부모를 찾을 때 까지 함수를 수행한다. 

 

union을 한다는 의미는 두 정점을 잇는다는 의미이고, cost가 현재 오름차순으로 정렬되어 있는 상황에서 작은 cost들 부터 vertex들을 union시키게 되면 결국 하나의 MST를 만들 수 있다는 점이 이 문제의 핵심이다.

 

union-find 문제를 많이 다뤄보지 않아서 조금 헷갈렸지만 개념에 대해서 잘 알고 있으면 충분히 풀 수 있던 문제였다.