C++ 프림(Prim) 알고리즘으로 MST 찾기

1. 프림(Prim) 알고리즘이란?

최소신장트리(MST, Minimum Spanning Tree) 를 찾기 위한 대표적인 알고리즘이다.

여기에서 최소신장트리란, 아래와 같이 모든 노드를 연결하는 부분 그래프 중, 가장 비용이 적은 것이다.

 

하늘색으로 표시한 경로가 최소신장트리이다.

여기에서 최소신장트리의 길이는 (4 + 8 + 16 + 20) = 48다.

 

최소신장트리를 찾는 알고리즘은 크루스칼(Kruskal) 과 프림(Prim) 이 대표적이다.

2. 동작원리

프림 알고리즘에서는 MST의 후보가 될 간선을 담을 우선순위 큐가 필요하다.

우선순위 큐는 최소의 비용을 가지는 경로가 우선순위를 갖게 한다.

(보통 Min Heap을 이용해서 구현한다.)

 

가장 먼저, 그래프에서 아무 노드나 선택하므로 가장 번호가 빠른 노드를 선택한다. 그리고, 선택된 노드에 연결된 모든 간선을 우선순위 큐에 넣는다.

 

그리고 이 중에서, 우선순위 큐에서 값을 하나 뺀다.

여기에서는 가장 위에 있는 가중치 4짜리 간선이 선택될 것이다.

이제 이 간선을 MST의 포함되는 간선으로 삼으며 아래와 같다.

 

이제 정점 1과 정점 2를 잇는 간선을 선택했으므로, 정점 2에 대해서 연결된 간선을 모두 우선순위 큐에 추가한다.

 

 

그리고 똑같이, 우선순위 큐에서 간선을 하나 빼서 MST 후보로 추가한다.

 

다음 차례는 위의 과정을 반복하므로 아래와 같은 상태가 된다.

 

하지만 아래에서처럼, 우선순위 큐에 있는 가중치 9 짜리의 간선을 선택하면 루프가 생기게 되므로 MST의 후보라고 생각하지 않는다. (빨간색으로 표시한 선)

 

그리고 다음 우선순위 큐의 원소를 추출하여 과정을 반복한다.

 

정점 3과 정점 4를 잇는 간선은 MST를 만들지 않으므로 이 간선 역시 MST의 후보가 되었다.

 

정점의 개수가 N 개일 때, 이렇게 우선순위 큐에서 뺀 간선을 MST 후보로 추가하는 과정을 간선이 N-1 개 만들어질 때까지 반복하면 MST가 만들어진다.

#include <iostream>

#define MAX_NODE	50000
#define MAX_EDGE	200000

using namespace std;

typedef struct Edge {
	int from;
	int to;
	int weight;
	Edge* next;
	Edge() { from = 0, to = 0, weight = 0, next = NULL; }
	Edge(int newFrom, int newTo, int newWeight) { from = newFrom, to = newTo, weight = newWeight, next = NULL; }
} Edge;

typedef struct List {
	Edge* head;
	List() { head = NULL; }
} List;

typedef struct Heap{
	Edge heap[MAX_EDGE];
	int heapSize;

	Heap() { heapSize = 0; }
} Heap;

// Variables
Heap h;
List* list[MAX_NODE];
int N, M;
bool included[MAX_NODE];

void listAdd(List* list, int from, int to, int weight)
{
	Edge* edge = new Edge(from, to, weight);
	Edge* head = list->head;

	if (list->head == NULL)
		list->head = edge;
	else 
	{
		while (head->next != NULL) head = head->next;
		head->next = edge;
	}
}

void heapInit(Heap* h) { h->heapSize = 0; }
void heapPush(Heap* h, Edge edge)
{
	int idx = ++(h->heapSize);

	while (idx != 1 && (h->heap[idx / 2].weight > edge.weight))
	{
		h->heap[idx] = h->heap[idx / 2];
		idx /= 2;
	}
	h->heap[idx] = edge;
}
Edge heapPop(Heap* h)
{
	Edge edge = h->heap[1];
	Edge temp = h->heap[(h->heapSize)--];
	int parent = 1;
	int child = 2;

	while (child <= h->heapSize)
	{
		if (child < h->heapSize && h->heap[child].weight > h->heap[child + 1].weight)
			child++;
		if (temp.weight <= h->heap[child].weight) break;

		h->heap[parent] = h->heap[child];
		parent = child;
		child *= 2;
	}
	h->heap[parent] = temp;
	return edge;
}

int prim()
{
	int answer = 0;

	for (int i = 0; i < MAX_NODE; i++) included[i] = false;
	included[0] = true;

	Edge edge;
	heapInit(&h);

	Edge* head = list[0]->head;
	while (head != NULL) {
		heapPush(&h, *head);
		head = head->next;
	}
	int numCheck = 0;
	while(numCheck != N - 1) {
		// Choose next edge
		edge = heapPop(&h);

		if (!included[edge.to])
		{
			answer += edge.weight;
			included[edge.to] = true;

			head = list[edge.to]->head;
			while (head != NULL) {
				if(!included[head->to])
					heapPush(&h, *head);
				head = head->next;
			}
			numCheck++;
		}
	}
	return answer;
}

int main(void)
{
	cin.tie(NULL);
	ios::sync_with_stdio(false);

	int answer = 0;

	cin >> N;

	for (int i = 0; i < N; i++)
		list[i] = new List();

	cin >> M;
	for (int i = 0; i < M; i++)
	{
		int from, to, weight;
		cin >> from >> to >> weight;
		listAdd(list[from - 1], from - 1, to - 1, weight);
		listAdd(list[to - 1], to - 1, from - 1, weight);
	}
	cout << prim() << '\n';

	for (int i = 0; i < N; i++) 
	{
		Edge* head = list[i]->head;
		while (head != NULL){
			Edge* tmp = head;
			head = head->next;
			delete tmp;
		}
		delete[] list[i];
	}
	return 0;
}

[알고리즘] 프림(Prim) 알고리즘으로 MST 찾기 (C++로 구현하기) (tistory.com)