BOJ - 최소 스패닝 트리(1197)

MST를 찾는 알고리즘 중 union-find를 이용하여 O(E*logE)에 찾을 수 있는 Kruskal algorithm을 이용하였다.

 

코드

#include<iostream>
#include<vector>
#include<algorithm>
using namespace std;
int V, E;
typedef struct graph
{
    int v1;
    int v2;
    int cost;
}graph;
vector<graph> v;
 
int getRoot(vector<int> parent, int idx)
{
    int i;
    for(i = idx; i != parent[i]; i = parent[i]);
    return i;
/*
    if (parent[idx] == idx) return idx;
    else return getRoot(parent, parent[idx]);
    */
}
 
bool find(vector<int> parent, int a, int b)
{
    a = getRoot(parent, a);
    b = getRoot(parent, b);
    if (a == b) return true;
    else return false;
}
 
void my_union(vector<int> &parent, int a, int b)
{
    a = getRoot(parent, a);
    b = getRoot(parent, b);
    if (a < b) parent[b] = a;
    else parent[a] = b;
}
 
bool comp(graph a, graph b)
{
    return a.cost < b.cost;
}
 
int main()
{
    cin.tie(NULL);
    ios_base::sync_with_stdio(false);
    cin >> V >> E;
    vector<int> parent(V + 1);
    for(int i=1; i<=V; i++)
        parent[i] = i;
    for(int i=0; i<E; i++)
    {
        graph g;
        cin >> g.v1 >> g.v2 >> g.cost;
        v.push_back(g);
    }
    sort(v.begin(), v.end(), comp);
    int sum = 0;
    for(int i=0; i<v.size(); i++)
    {
        if (!find(parent, v[i].v1, v[i].v2))
        {
            my_union(parent, v[i].v1, v[i].v2);
            sum += v[i].cost;
        }
    }
    printf("%d\n", sum);
    return (0);
}

그런데, 주의해야 할 점이 있었다. 기존에는 getRoot를 찾을 때 재귀로 찾는 방법을 선택하였는데 시간초과가 떠버린 것이다. 

그래서 그냥 return하는 것이 아니라 parent[idx]의 값을 수정하면서 return하게 하는 path-compression을 이용하였는데도 시간초과가 떠서, 위와 같이 for문으로 구현하였더니 시간초과를 피할 수 있었다. 

 

추가)

int getRoot(vector<int> &parent, int idx)
{
    /*
    int i;
    for(i = idx; i != parent[i]; i = parent[i]);
    return i;*/
    if (parent[idx] == idx) return idx;
    else
    {
        int root = getRoot(parent, parent[idx]);
        parent[idx] = root;
        return parent[idx];
    }
}