문제 1197. 최소 스패닝 트리(에센셜 5)

문제 링크

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

 

난이도 : 골드 4

 

문제 요약 설명

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

최소 스패닝 트리는 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다. 그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 

입력

1번째 줄: 정점의 개수 V 간선의 개수 E

다음 E개의 줄: 세 정수 A, B, C ( A번 정점과 B번 정점이 가중치 C인 간선으로 연결)

 

입력 제한

정점의 개수 V(1 ≤ V ≤ 10,000)

간선의 개수 E(1 ≤ E ≤ 100,000)

0 <= |가중치| <  1,000,000

최소 스패닝 트리의 가중치는 int 범위 내의 데이터만 입력으로 주어진다.

입력 예제

// input
3 3
1 2 1
2 3 2
1 3 3

// ans
3

 

풀이 방식

프림 알고리즘(Prim algorithm)이나 크루스칼 알고리즘(kruskal algorithm)을 사용하면 된다. 워낙 유명한 문제이기 때문에 설명은 생략하겠다.

 

정답 코드 

Prim algorithm

#include <bits/stdc++.h>

#define fastio cin.tie(0)->sync_with_stdio(0)

using namespace std;

int V, E;

struct edge_data
{
    int node_no;
    int weight;
};

struct cmp
{
    bool operator()(const edge_data& a, const edge_data& b){
        return a.weight > b.weight;
    }
};

vector<vector<edge_data>> graph;
vector<bool> visited;

int prim_algorithm(){
    int edge_sum = 0;

    priority_queue<edge_data, vector<edge_data>, cmp> pq;

    int next_node_no = 1;
    for(int x = 0; x < V-1; x++){
       
        visited[next_node_no] = true;

        for(edge_data ed : graph[next_node_no]){
			if(!visited[ed.node_no]) pq.push(ed);     
        }

        while(!pq.empty()){
            edge_data cur_edge = pq.top(); pq.pop();
            if(!visited[cur_edge.node_no]){
                edge_sum += cur_edge.weight;
                next_node_no = cur_edge.node_no;
                break;
            }
        }
    }
      
    return edge_sum;
}

int main() {
	fastio;
    cin >> V >> E;
   
    graph.resize(V+1);
    visited.resize(V+1);
    
    int node1, node2, weight;
    for(int i = 0; i < E; i++){
        cin >> node1 >> node2 >> weight;
        graph[node1].push_back({node2, weight});
        graph[node2].push_back({node1, weight});
    }   
    
    int ans = prim_algorithm();

    cout << ans << '\n';

    return 0;
}

kruskal algorithm

#include <bits/stdc++.h>

#define fastio cin.tie(0)->sync_with_stdio(0)

using namespace std;

int V, E;

vector<int> root;

struct edge_data
{
    int node1;
    int node2;
    int weight;

};

struct cmp
{
    bool operator()(const edge_data& a, const edge_data& b){
        return a.weight > b.weight;
    }
};

priority_queue<edge_data, vector<edge_data>, cmp> pq;

/* find(x): 재귀 이용 */
int find(int x) {
    // 루트 노드는 부모 노드 번호로 자기 자신을 가진다.
    if (root[x] == x) {
        return x;
    } else {
        // 각 노드의 부모 노드를 찾아 올라간다.
        return find(root[x]);
    }
}

/* union(x, y) */
void union_set(int x, int y){
    // 각 원소가 속한 트리의 루트 노드를 찾는다.
    x = find(x);
    y = find(y);

    if(x < y){
        root[y] = x;
    }
    else{
        root[x] = y;
    }
}

int kruscal_algorithm(){
    int edge_sum = 0;

    while(!pq.empty()){
        edge_data cur_edge = pq.top(); pq.pop();

        //cout << "ssssssssss "<< cur_edge.node1 << " " << cur_edge.node2 << " " << cur_edge.weight << '\n';
        
        if (find(cur_edge.node1) != find(cur_edge.node2)){
            //cout << "AADD" << '\n';
            union_set(cur_edge.node1, cur_edge.node2);
            edge_sum += cur_edge.weight;
        }

    }

    return edge_sum;
}

int main() {
	fastio;
    cin >> V >> E;
  
    root.resize(V+1);

    for(int i = 0; i < root.size(); i++){
        root[i] = i;
    }

    if(V == 1){
        cout << 0 << '\n';
        return 0;
    }

    int node1, node2, weight;
    for(int i = 0; i < E; i++){
        cin >> node1 >> node2 >> weight;
        pq.push({node1, node2, weight});
    }   

    // while(!pq.empty()){
    //     cout << pq.top().node1 << " " << pq.top().node2 << " " << pq.top().weight << '\n';
    //     pq.pop();
    // }
    
    int ans = E == 1 ? weight : kruscal_algorithm();

    cout << ans << '\n';

    return 0;
}

더 좋은 코드?

 

참고