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

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

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

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

 

 

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

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

 

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

오늘은 그 중에서도 프림 알고리즘에 대해서 알아보겠습니다.

 

 

 

동작원리

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

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

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

 

가장 먼저, 그래프에서 아무 노드나 선택합니다.

아무 노드라고 했으므로, 여기에서는 가장 번호가 빠른 노드를 선택하겠습니다.

그리고, 선택된 노드에 연결된 모든 간선을 우선순위 큐에 넣습니다.

 

 

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

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

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

 

이제 정점 1과 정점 2를 잇는 간선을 선택했으므로,

정점 2에 대해서 연결된 간선을 모두 우선순위 큐에 추가합니다.

 

 

그리고 똑같이, 우선순위 큐에서 간선을 하나 빼서

MST 후보로 추가합니다.

 

 

다음 차례입니다. 위의 과정을 반복하므로

아래와 같은 상태가 됩니다.

 

하지만 아래에서처럼, 우선순위 큐에 있는

가중치 9 짜리의 간선을 선택하면 루프가 생기게 됩니다.

그러므로 MST의 후보라고 생각하지 않습니다. (빨간색으로 표시한 선)

 

 

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

 

정점 3과 정점 4를 잇는 간선은 MST를 만들지 않네요.

이 간선 역시 MST의 후보가 되었습니다.

 

정점의 개수가 N 개일 때,

이렇게 우선순위 큐에서 뺀 간선을 MST 후보로 추가하는 과정을

간선이 N-1 개 만들어질 때까지 반복하면 MST가 만들어집니다.

 

 

 

구현 (C++)

#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;
}