Prim 알고리즘 자바로 구현 (백준 1197번)

2021.08.30 - [알고리즘] - Minimum Spanning Tree를 구현하는 Prim 알고리즘에 대해

Minimum Spanning Tree를 구현하는 Prim 알고리즘에 대해

저번 포스팅에서 프림 알고리즘에 대해 다루었다.

그래프로 설명하면 어렵지 않지만 이를 자바 코드로는 어떻게 구현할 수 있을까?

참고로 Minimum spanning tree를 코드로 구현하는 문제를 풀고 싶다면 백준 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보다 작거나 같은 데이터만 입력으로 주어진다.

출력

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

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

1197번: 최소 스패닝 트리

---

항상 연결된 간선 중 가장 비용이 작은 간선을 택하면서 vertex를 선택하여 나가는 것이 prim 알고리즘의 핵심이므로 우리는 PriorityQueue를 사용할 것이다. 그리고 PriorityQueue를 사용하기 위해서는 Comparable이 구현되어 있어야 하므로 간선의 weight를 비교 기준으로 삼는 class를 하나 생성한다. 

 

import java.util.*;
import java.io.*;
class Node implements Comparable<Node> {
    int index;
    int weight;
    Node(int index, int weight){
        this.index=index;
        this.weight=weight;
    }
    @Override
    public int compareTo(Node o){
        return Integer.compare(this.weight, o.weight);
    }
}
public class Main{
    public static void main(String[] args) throws IOException{
        BufferedReader br=new BufferedReader(new InputStreamReader(System.in));
        String a1[]=br.readLine().split(" ");
        int v=Integer.parseInt(a1[0]);
        int e=Integer.parseInt(a1[1]);
        ArrayList<Node> arr[]=new ArrayList[v+1];
        for(int i=1; i<=v; i++){
            arr[i]=new ArrayList<>();
        }
        for(int i=0; i<e; i++){
            String infos[]=br.readLine().split(" ");
            arr[Integer.parseInt(infos[0])].add(new Node(Integer.parseInt(infos[1]), Integer.parseInt(infos[2])));
            arr[Integer.parseInt(infos[1])].add(new Node(Integer.parseInt(infos[0]), Integer.parseInt(infos[2])));       
        }
         
        //Prim 알고리즘 시작
        boolean visited[]=new boolean[v+1];
        PriorityQueue<Node> pq=new PriorityQueue<>();
        pq.add(new Node(1, 0));
        long result=0;
        int count=0;
        while(!pq.isEmpty()){
            Node current=pq.poll();
            if (visited[current.index]){
                continue;
            }
            visited[current.index]=true;
            result+=current.weight;
            if (++count==v){
                System.out.println(result);
                break;
            }
            for(Node next: arr[current.index]){
                   pq.add(next);
               
            }
        }
        //Prim end//
    }
}

전체 코드는 위와 같다. Prim 알고리즘 부분은 주석으로 체크해두었으니 참고하면 될 것 같다. 

필요한 배열은 visited라는 배열과 ArrayList의 집합이다. Prim 알고리즘에서는 한 번 방문한 vertex는 다시 방문하지 않으므로 이 부분의 배열 체크가 필요하다. 또한 연결된 vertex를 저장하기 위해 ArrayList의 배열이 필요하다. 

 

처음 시작 노드는 아무거나 설정해도 됨으로 1번 vertex를 넣어주었고 이때는 간선이 없으므로 간선의 weight는 0으로 설정해두었다.

 

그리고 count에는 선택된 vertex의 개수를 저장하고 result에는 간선의 총합을 저장한다. 전체 vertex의 개수와 count의 개수가 같아지면 모든 vertex를 다 찾은 것이므로 종료하면 된다.

 

vertex가 선택될 때마다 그와 연결된 모든 vertex를 prioirtyQueue에 넣어둔다. 그렇기 때문에 visited를 여기서 처리하는 것이 아니라 이후 priorityQueue에서 poll() 할때가 vertex가 방문되는 순서라고 할 수 있다. 

 

PriorityQueue를 사용하기 때문에 vertex를 poll할 때마다 최소 간선이 꺼내진다. 따라서 우리가 매번 sorting을 하지 않아도 Prim 알고리즘이 자동으로 구현되는 효과를 볼 수 있다.