[백준 알고리즘] 1197번: 최소 스패닝 트리

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

문제

문제
그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력
첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력
첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

풀이

크루스칼 알고리즘을 그대로 구현했더니 쉽게 풀렸다. 그런데 다른 풀이들을 보니 역시 크루스칼 알고리즘은 유니온 파인드의, 특히 파인드 연산 속도에 큰 영향을 받는 것 같다. 컨테이너를 어떤 걸 쓸지 고민이 좀 됐는데, 그냥 벡터 하나를 사용해도 되지만, 좀 더 의미를 부여하기 위해서 복잡해도 중첩 pair를 벡터안에 넣어서 사용했다. 정점 두개가 한쌍이 되고, 이 두개의 정점을 잇는 간선의 가중치가 또 다른 짝이 되도록 말이다. 유니온 파인드의 merge 연산에서 사이클이 발견되면 false 가 리턴되게 해서 이 경우에는 간선 가중치를 결과 변수에 더하지 않고, 성공적으로 merge 연산이 수행됐을 때만 가중치를 누적합 시켰다.

코드

#include <cstdio>
#include <vector>
#include <algorithm>

using namespace std;

int parent[10001];
int level[10001];
vector<pair<pair<int, int>, int>> edges;

bool comp (pair<pair<int, int>, int> a, pair<pair<int, int>, int> b){
    return a.second < b.second;
}

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

bool merge(int a, int b){
    a = findRoot(a);
    b = findRoot(b);

    if (a == b) return false;
    if (level[a] > level[b]) parent[a] = b;
    else parent[b] = a;

    if (level[a] == level[b]) level[a]++;
    return true;
}

int main (){
    int v, e;
    int result = 0;
    scanf("%d %d", &v, &e);

    for (int i = 0 ; i <= v ; i++){
        parent[i] = i;
    }

    for (int i = 0 ; i < e ; i++){
        int a, b, c;
        scanf("%d %d %d", &a, &b, &c);
        edges.push_back(make_pair(make_pair(a, b), c));
    }

    sort(edges.begin(), edges.end(), comp);

    for (auto i : edges){
        if(merge(i.first.first, i.first.second)) {
            result += i.second;
        }
    }

    printf("%d", result);
}

Written by@전여훈 (Click Me!)
고민이 담긴 코드를 만들자, 고민하기 위해 공부하자.