[백준] 1197번 최소 스패닝 트리 (JAVA)

문제

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

1197번: 최소 스패닝 트리

 

설명

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

이번 문제에서는 프림 알고리즘을 통해 최소 신장 트리를 구할 것이다.

먼저 간선의 정보를 저장할 Edge 클래스를 만든다.

우선순위 큐에서 가중치를 기준으로 오름차순으로 정렬하기 위해 Comparable 인터페이스를 이용한다.

class Edge implements Comparable<Edge> {
    int dest;
    int cost;

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

    @Override
    public int compareTo(Edge edge) {
        return this.cost - edge.cost;
    }
}

각 정점의 방문 여부를 확인하기 위한 visited 배열을 만든다.

우선순위 큐에 시작 정점인 1을 넣는다.

큐가 비어있을 때까지 큐에서 하나씩 꺼내어 방문한 정점이라면 건너뛴다.

방문한 정점이 아니라면 방문 처리를 하고 minCost에 그 정점의 가중치를 더한다.

그 다음 해당 정점에 연결되어 있는 모든 간선을 탐색하여 방문하지 않은 간선을 큐에 넣는다.

static void prim() {
    boolean[] visited = new boolean[V + 1];

    pq.offer(new Edge(1, 0));

    int minCost = 0;

    while (!pq.isEmpty()) {
        Edge edge = pq.poll();
        int dest = edge.dest;
        int cost = edge.cost;

        // 방문 했다면 건너뜀
        if (visited[dest])
            continue;

        visited[dest] = true;
        minCost += cost;

        // 해당 정점에 연결되어 있는 모든 간선 탐색
        for (Edge e : graph.get(dest)) {
            if (!visited[e.dest]) {
                pq.add(e);
            }
        }
    }

    System.out.println(minCost);
}

 

전체 코드

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

public class Main {

    static int V, E;
    static List<List<Edge>> graph;
    static PriorityQueue<Edge> pq;

    static class Edge implements Comparable<Edge> {
        int dest;
        int cost;

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

        @Override
        public int compareTo(Edge edge) {
            return this.cost - edge.cost;
        }
    }

    static void prim() {
        boolean[] visited = new boolean[V + 1];

        pq.offer(new Edge(1, 0));

        int minCost = 0;

        while (!pq.isEmpty()) {
            Edge edge = pq.poll();
            int dest = edge.dest;
            int cost = edge.cost;

            // 방문 했다면 건너뜀
            if (visited[dest])
                continue;

            visited[dest] = true;
            minCost += cost;

            // 해당 정점에 연결되어 있는 모든 간선 탐색
            for (Edge e : graph.get(dest)) {
                if (!visited[e.dest]) {
                    pq.add(e);
                }
            }
        }

        System.out.println(minCost);
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

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

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

        for (int i = 0; i < E; i++) {
            st = new StringTokenizer(br.readLine());
            int A = Integer.parseInt(st.nextToken());
            int B = Integer.parseInt(st.nextToken());
            int C = Integer.parseInt(st.nextToken());

            graph.get(A).add(new Edge(B, C));
            graph.get(B).add(new Edge(A, C));
        }

        pq = new PriorityQueue<>();

        prim();
    }
}