[백준] 최소 스패닝 트리(1197) Java, Prim 알고리즘

문제 링크

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

 

문제 요약

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

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

 

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다.

다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다.

이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다.

C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

 

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다.

최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

 

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

 

접근법

단순히 주어진 그래프에서 MST를 구하는 문제이므로, Kruskal 알고리즘이나 Prim 알고리즘을 구현하기만 하면 된다.

 

맨 처음에는 Kruskal 알고리즘으로 문제를 풀어서 맞았고, 공부를 위해 Prim 알고리즘으로 다시 구현해봤다.

Kruskal 알고리즘이 궁금하다면, 여기를 클릭하시라.

 

Prim 알고리즘은 Kruskal과 달리 edge가 기준이 아니라 node가 기준이다.

Kruskal에서는 MST가 완성될 때 까지 cycle의 여부를 union-find로 체크해가며 edge를 greedy하게 추가했다면,

Prim에서는 greedy하게 node를 하나씩 추가해가며 모든 node를 포함할 때 까지 tree를 키워나가는 형태이다.

근데 실제 구현에서는 Prim 알고리즘도 마치 edge를 추가하는 형태처럼 구현이 되어서, 코드 측면에서는 kruskal 보다 이해가 어렵다는 단점이 있다.

 

그리고 edge가 많은 dense graph가 아닌 이상은 Kruskal과 Prim이 성능면에서 크게 차이가 나지 않는다.

Kruskal은 O(ElogE)이고, Prim은 O(ElogV)이기 때문.

그래서 그냥 왠만하면 Kruskal을 쓰는게 편할 듯 하다.

 

아무튼 아래와 같은 그래프에서 Prim으로 MST를 찾는다고 해보자.

Prim 알고리즘은 다음과 같은 과정으로 진행된다.

1. 임의의 node 하나를 시작 node로 정한다. 시작 node가 MST에 추가되었다고 표시한다. MST는 모든 node를 포함하기 때문에 어떤 node를 선택해도 상관 없다.
2. 시작 node와 연결된 모든 edge(혹은 node로 봐도 된다)를 PriorityQueue에 넣는다.
3. PriorityQueue에서 edge를 poll한다.
4. Poll 한 edge(의 다른 한 쪽 node)가 이미 MST에 포함되어 있으면 edge를 새로 poll 한다.
5. MST에 포함되어 있지 않으면 MST에 추가되었다고 표시하고, 해당 node와 연결되어 있는 edge(node)를 PriorityQueue에 넣는다. 이 때 이미 MST에 포함되어 있는 edge(node)는 넣지 않는다. 
6. PriorityQueue가 비거나 V - 1개의 edge를 MST에 포함시킬 때 까지 2-5를 반복한다.

 

이 때 edge와 node가 동일하게 취급되는 이유는, 실제 cost는 edge가 갖고 있기 때문이다. 그림으로 보면,

같은 색의 edge와 node가 한 묶음이라고 볼 수 있다

1번 node에서 연결된 edge(node)를 PriorityQueue에 추가할 때, 1번 node에서 2번 node로 향하는 cost 1의 edge와 2번 node를 한 묶음으로, 또 1번 node에서 3번 node로 향하는 cost 7의 edge와 3번 node를 한 묶음으로 쳐서 추가하는 느낌으로 보면 된다.

 

그러니까 Prim 알고리즘은 정확히는 node를 MST에 하나씩 추가해가며 전체 node를 다 추가할 때 까지 반복하는 알고리즘이지만, 실제 cost를 갖고 있는 것은 edge이기 때문에 마치 edge와 node 쌍을 추가하는 것처럼 구현된다는 것이다.

 

사실 말로 아무리 해봤자 이해가 어렵고 코드로 보면 편하다.

 

 

시간 복잡도

Prim 알고리즘은 O(ElogV)이다.

 

 

공간 복잡도

정확히 계산이 어렵다...

 

 

풀면서 놓쳤던 점

딱히 없다.

 

 

이 문제에서 얻어갈 점

Prim이 Kruskal보다 구현이 까다롭다는 사실을 배웠다.

 

 

코드

import java.util.*;
import java.io.*;

class Main
{
	public static void main(String[] args) throws Exception
	{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;

		int V, E;
		st = new StringTokenizer(br.readLine());
		V = Integer.parseInt(st.nextToken());
		E = Integer.parseInt(st.nextToken());

		ArrayList<ArrayList<Edge>> graph = new ArrayList<>(V + 1);
		for (int i = 0; i <= V; i++)
			graph.add(new ArrayList<Edge>());

		int node1, node2;
		long cost;
		for (int i = 1; i <= E; i++)
		{
			st = new StringTokenizer(br.readLine());
			node1 = Integer.parseInt(st.nextToken());
			node2 = Integer.parseInt(st.nextToken());
			cost = Long.parseLong(st.nextToken());

			//node를 기준으로 동작하는 알고리즘이기 때문에
			//각 node에 인접한 node와 그 cost를 추가해주는 방식
			graph.get(node1).add(new Edge(node2, cost));
			graph.get(node2).add(new Edge(node1, cost));
		}

		//MST에 추가된 node는 true
		boolean[] selected = new boolean[V + 1];

		PriorityQueue<Edge> edges = new PriorityQueue<Edge>();
		//1번 node에서 시작하기 위해
		//존재하지 않는 node에서 1번 node로 이어진 cost 0의edge가 있다고 가정하고
		//PriorityQueue에 추가함
		edges.add(new Edge(1, 0));
		
		cost = 0;
		int cnt = 0;
		Edge e;
		//cnt == V -1 이 아닌 cnt == V 일 때 끝나는 이유는
		//가상의 edge인 Edge(1,0)을 추가했기 때문!
		while (!edges.isEmpty() && cnt < V)
		{
			e = edges.poll();

			if (selected[e.node])
				continue;

			cost += e.cost;
			selected[e.node] = true;
			cnt++;

			//인접한 node 중 MST에 추가되지 않은 node를 pq에 넣는다.
			for (Edge adj : graph.get(e.node))
			{
				if (selected[adj.node] == false)
					edges.add(adj);
			}
		}

		System.out.println(cost);

		br.close();
	}
}

class Edge implements Comparable<Edge>
{
	int node;
	long cost;

	public Edge(int node, long cost)
	{
		this.node = node;
		this.cost = cost;
	}

	public int compareTo(Edge other)
	{
		if (this.cost > other.cost)
			return 1;
		else if (this.cost < other.cost)
			return -1;
		else
			return 0;
	}
}