[알고리즘] 프림 알고리즘 (Prim's Algorithm), C++

프림 알고리즘

프림 알고리즘은 최소 신장 트리를 구하는 방법 중 하나입니다.

 

하나의 정점을 시작으로 MST 정점 집합을 점차 넓혀가는 방법입니다.

1. 시작 정점을 집합에 추가합니다.
2. 집합 내부의 정점들과 인접한 정점들 중 간선의 가중치가 최소로 연결된 정점을 집합에 추가합니다.
3. 집합 내부의 정점 개수가 n개가 될 때가지 2번의 과정을 반복합니다.

출처: https://gmlwjd9405.github.io/2018/08/30/algorithm-prim-mst.html

 

 

구현 방법

집합에 포함된 정점들과 인접한 간선들 중 최소 가중치를 가지는 간선을 효과적으로 찾기 위해서는 우선순위 큐를 활용하는 것이 좋습니다.
만약 일반 배열을 사용한다면 매번 간선들의 최소 값을 찾기 위해 $n^2$의 시간을 사용하게 됩니다.

 

1) 초기 상태

시작 정점을 0번으로 하기 때문에 우선순위 큐에는 (0, 0)을 삽입합니다.

 

2) (0, 0) 추출

0번 노드와 연결된 1, 2, 4번 노드와 그 가중치가 우선순위 큐에 삽입됩니다. 
visited[0] = true 변경해 줍니다.

 

3) (1,2) 추출

1번 노드와 인접한 0, 3 번 노드와 그 가중치를 우선순위 큐에 삽입해야 하는데, 이때 0번 노드의 visitied 값이 true 이기 때문에 0번 노드와 연결된 간선은 큐에 삽입하지 않습니다.

선택한 간선의 가중치를 전체 가중치 합에 누적합니다.

 

 

4) (3, 2) 추출

 

5) (4, 3) 추출

 

 

 

6) (2, 3) 추출

MST 정점 집합의 원소 수가 5개가 되었기 때문에 반복을 종료하고 total cost를 출력합니다.

 

 

추가)

위의 예시에서는 보여지지 않았지만,
우선순위 큐에서 나온 current 정점이 이미 방분된(visited[cur] == true) 인 경우가 발생할 수 있습니다.

큐에서 나온 정점이 이미 방문된 상태라면 해당 경우를 넘기고(continue) 큐에서 다음 정점을 처리해야 합니다.

 

관련 문제

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

1197번: 최소 스패닝 트리

 

코드

#include <iostream>
#include <vector>
#include <queue>

using namespace std;

struct Edge
{
    int to;
    int weight;

    bool operator>(Edge& input)
    {
        return this->weight > input.weight;
    }
};

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

int v, e;
vector<vector<Edge>> graph;

int main()
{
    cin.tie(0);
    ios_base::sync_with_stdio(false);

    cin >> v >> e;
    graph.resize(v+1);
    for (int i=0; i<e; ++i)
    {
        int a, b, weight;
        cin >> a >> b >> weight;
        graph[a].push_back({b, weight});
        graph[b].push_back({a, weight});
    }
    
    priority_queue<Edge, vector<Edge>, greater<Edge>> pq;
    vector<bool> visited(v+1, false);
    int totalCost = 0;
    vector<int> MST;
    pq.push({1, 0});
    
    while(!pq.empty())
    {
        int cur = pq.top().to;
        int curWeight = pq.top().weight;
        pq.pop();

        if (visited[cur])
            continue;

        MST.push_back(cur);
        totalCost += curWeight;
        visited[cur] = true;

        if (MST.size() == v)
            break;

        for (Edge e : graph[cur])
        {
            int next = e.to;
            int nextWeight = e.weight;

            if (visited[next])
                continue;

            pq.push(e);
        }

    }

    cout << totalCost << endl;
    return 0;
}