백준 1197 최소 스패닝 트리 (자바)

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

1197번: 최소 스패닝 트리

 

최소 스패닝 트리(MST)에 대한 문제다.

최소 스패닝 트리란?

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

MST개념을 공부하면 가장 처음 접하게 되는 문제다. (골드 4지만 방법을 알면 실버급이라고 생각한다.)

 

MST를 구하는 대표적인 방법으로는 크루스칼, 프림 2가지가 있다.

나는 이중에서 프림(우선순위큐)을 사용한 방법으로 문제를 풀었다.

 

대부분 인터넷을 찾아보면 프림은 간선의 개수가 많을 때 유리하고, 크루스칼은 간선의 개수가 적을 때 유리하다고 한다.

하지만 아직 세부적인 접근은 힘들다고 판단해서 우선 프림을 사용하는 방법을 거의 외웠다.

 

입력 

3 3
1 2 1
2 3 2
1 3 3

입력으로 첫 줄에는 정점, 간선이 들어온다

이후에는 A정점, B정점, 가중치 순서대로 들어온다.

 

출력

정답으로 3이 출력된다.

 

입력 예시 그림

입력을 그림으로 표현을 해봤다.

1에서 3으로 가는 방법은 총 2가지가 있다.

1) 1-> 2-> 3

2) 1-> 3

이 2가지 방법 모두 가중치는 최소 3이 들어간다. 그래서 정답으로 3이 나온 것이다.

 

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Scanner;

public class Main {
	public static class Edge implements Comparable<Edge>{
		int ed;
		int w;
		public Edge( int ed, int w) {
			this.ed = ed;
			this.w = w;
		}
		@Override
		public int compareTo(Edge o) {
			return this.w - o.w;
		}
		
	}
	
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int V = sc.nextInt();
		int E = sc.nextInt();
		
		List<Edge>[] list = new ArrayList[V+1];
		for(int i=0; i<V+1; i++) {
			list[i] = new ArrayList<>();
		}
		
		for(int i=0; i<E; i++) {
			int st = sc.nextInt();
			int ed = sc.nextInt();
			int w = sc.nextInt();
			
			list[st].add(new Edge(ed, w)); // 무방향이기때문
			list[ed].add(new Edge(st, w));
		}
		
		PriorityQueue<Edge> pq = new PriorityQueue<>();
		boolean[] visited = new boolean[V+1];
		pq.addAll(list[1]);
		visited[1] = true;
		
		int pick = 1;
		long sum = 0;
		while(pick!=V) {
			Edge e = pq.poll();
			if(visited[e.ed]) continue;
			
			visited[e.ed] = true;
			sum += e.w;
			pq.addAll(list[e.ed]);
			pick++;
		}
		
		System.out.println(sum);
	}
}

우선 Edge라는 클래스를 만들었다.

처음에는 무슨 소리인지 이해 안됐는데 이런 방법을 다들 많이 사용하는 것 같았다.

public static class Edge implements Comparable<Edge>{
		int ed;
		int w;
		public Edge( int ed, int w) {
			this.ed = ed;
			this.w = w;
		}
		@Override
		public int compareTo(Edge o) {
			return this.w - o.w;
		}
		
	}

우선 클래스를 먼저 살펴보면 도착정점과 가중치를 받았다.

가중치를 compareTo로 오름차순 정렬을 했는데 그 이유는 최소를 가지는 경로를 찾아야 하기 때문에

가중치를 오름차순으로 정렬해서 최소부터 빼기 위해 사용했다.

 

public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int V = sc.nextInt();
		int E = sc.nextInt();
		
		List<Edge>[] list = new ArrayList[V+1];
		for(int i=0; i<V+1; i++) {
			list[i] = new ArrayList<>();
		}
		
		for(int i=0; i<E; i++) {
			int st = sc.nextInt();
			int ed = sc.nextInt();
			int w = sc.nextInt();
			
			list[st].add(new Edge(ed, w)); // 무방향이기때문
			list[ed].add(new Edge(st, w));
		}

그 다음 정점과 간선을 입력받는다.

이후에 Edge타입을 가진 리스트를 담을 수 있는 배열을 생성했다.

for문을 통해서 배열에 리스트를 생성해 줬다.

 

간선의 개수만큼 입력을 받는다. 보기 편하게 시작 정점, 도착 정점, 가중치로 나눠서 받았다.

list[st].add(new Edge(ed, w)); 이 부분은 리스트의 시작 정점을 배열 번호로 두고 도착 정점과 가중치를 Edge객체 안에 집어넣은 것이다.

 

만약에 1,2,3 / 1,3,4 이렇게 2개의 입력이 들어왔다고 하면

list[1]번에 도착 정점이 2번이고 가중치가 3인 정보, 도착 정점이 3번이고 가중치가 4인 정보가 들어있는 것이다.

 

무방향 그래프이기 때문에 반대로 뒤집어서도 리스트에 넣어줬다.

 

PriorityQueue<Edge> pq = new PriorityQueue<>();
		boolean[] visited = new boolean[V+1];
		pq.addAll(list[1]);
		visited[1] = true;
		
		int pick = 1;
		long sum = 0;

 

그다음 우선순위큐를 사용하기 때문에 우선순위큐를 먼저 생성해 주었다.

방문을 했는지 확인하기 위해 boolean으로 visited배열을 생성하고

큐에 addAll을 써서 우선 1번 정점에 들어있는 리스트를 추가하였다.

 

기존에는 보통 pq.add를 많이 사용했다. 하지만 이번 문제에서는 1번 정점에서 시작하는 간선이 여러 개 존재할 수가 있다.

1,3 / 1,4 / 1,5 이렇게 1번에서 출발해서 3,4,5번으로 도착하는 간선이 존재한다고 생각해 보자.

리스트에 3개의 정보가 담겨있기 때문에 이 모든 정보를 pq에 넣으려면 for문을 돌려서 넣어야 했다.

 

하지만 더 효율적이고 간단하게 집어넣기 위해서 addAll을 사용하였다.

 

그리고 1번은 이미 들렸으니 visited[1]번을 true로 만들어준다.

 

int pick은 내가 선택한 정점의 개수이다. 하나를 이미 넣었으니 1로 초기화를 해주었다.

long sum은 내가 더한 모든 가중치의 값이다.

 

while(pick!=V) {
			Edge e = pq.poll();
			if(visited[e.ed]) continue;
			
			visited[e.ed] = true;
			sum += e.w;
			pq.addAll(list[e.ed]);
			pick++;
		}

while문을 돌리는데 조건은 pick!=V때 동작하게 돌린다.

pick은 내가 선택한 정점의 개수인데 모든 정점을 골랐으면 while문이 멈춰야 하기 때문이다.

 

Edge타입의 변수 e를 만들어서 pq에 있는 데이터를 우선 꺼낸다.(여기서는 1번 정점에서 가중치가 가장 낮은 리스트 값이 꺼내지게 된다)

만약에 도착 정점이 이미 true라면 건너뛴다.

 

아니라면 true로 만들어서 방문했다고 표시하고

sum에다가 해당하는 가중치를 더해준다.

그리고 pq에 다시 도착 정점의 정보를 담아준다.

하나 정점을 방문했으니 pick을 하나 더해준다.

 

이 과정을 거치면 모든 정점을 연결하면서 최소 가중치를 가지는 값을 도출하게 된다.

 

완전 처음 접했던 내용이라 거의 코드를 외우기만 했었는데 블로그에 복기를 하면서 내용을 조금씩 정리해 가는 느낌이다.