10-5. [C/C++] MST를 위한 Prim 알고리즘

Prim 알고리즘

Prim 알고리즘은 하나의 정점에서부터 시작해 트리를 단계적으로 확장해나가는 방법이다.

 

처음에는 시작 정점만이 MST에 포함된다.

다음으로 지금까지 만들어진 MST에 인접한 정점들 중에서 간선의 가중치가 가장 작은 정점을 선택하여 트리를 확장한다.

이 과정을 트리가 (n - 1)개의 간선을 가질 때까지 계속한다.

Prim()

(1) 그래프에서 시작 정점을 선택해 초기 MST를 만든다.
(2) 현재 MST의 정점들과 인접한 정점들 중에서 간선의 가중치가 가장 작은 정점 v를 선택한다.
(3) 이 정점 v와 이때의 간선을 MST에 추가한다.
(4) 모든 정점이 삽입될 때 까지 2번으로 이동한다.

Kruskal 알고리즘이 간선을 중심으로 하는 알고리즘인데 비해 Prim 알고리즘은 정점을 중심으로 한다.

---

Prim 알고리즘의 구현

Prim 알고리즘 자체는 크게 복잡하지는 않지만 구현하기 위해 고려해야 할 것들이 있다.

 

첫 번째는 현재 MST에 인접한 정점들 중에서 가장 비용이 적게 드는 간선을 찾는 과정이다.

이것을 위해 우선 순위 큐를 사용할 것이다.

가중치는 기준으로 우선 순위 큐(최소 힙)를 사용하면 큐에서 가중치가 가장 작은 간선을 선택해 줄 것이다.

 

두 번째는 MST가 사이클을 형성하게 되는 것을 막아야 하는 것이다.

Prim 알고리즘은 간선을 중심으로 하는 Kruskal 알고리즘과 달리 정점을 중심으로 한다.

그렇기 때문에 집합이 형성되는 지 확인하는 것이 아니고, 정점이 이미 MST에 포함되어 있는지 검사한다.

이를 위해 bool visited[] 배열을 사용할 것이다.

---

코드로 구현

1197번: 최소 스패닝 트리

이 문제를 Prim 알고리즘을 이용해 해결해 봤다.

// prim.cpp
#include <iostream>
#include <vector>
#include <queue>
using namespace std;
#define SIZE 10001

int V, E;
vector<pair<int, int>> graph[SIZE];
priority_queue<pair<int, int>> pq;
bool visited[SIZE];
long long ans;

void prim() {
    pq.push({0, 1});

    while (!pq.empty()) {
        int curWeight = -pq.top().first;
        int curNode = pq.top().second;
        pq.pop();

        if (visited[curNode]) continue;

        visited[curNode] = true;
        ans += curWeight;

        for (int i = 0; i < graph[curNode].size(); i++) {
            int nextNode = graph[curNode][i].first;
            int nextWeight = graph[curNode][i].second;

            pq.push({-nextWeight, nextNode});
        }
    }
}

int main() {
    cin >> V >> E;
    for (int i = 0; i < E; i++) {
        int u, v, weight;
        cin >> u >> v >> weight;
        graph[u].push_back({v, weight});
        graph[v].push_back({u, weight});
    }

    prim();
    cout << ans << endl;

    return 0;
}

위 코드는 아래 이미지를 클릭해 깃허브에서도 확인할 수 있다.