[SWEA] 4006. 고속도로 건설 2 (with JAVA)

문제 4006

SW Expert Academy

 

[ 접근 과정 ]

문제에서 이미 '최소 비용' 이라는 말이 나와있어서, 가중치를 계산하는 문제라는 것을 알 수 있었습니다. '최소 비용으로 모든 도시 간을 이동할 수 있게' 라는 부분에서 크루스칼 알고리즘을 사용하면 좋을 것 같다는 생각을 했습니다. 여기서의 간선은 도로이기 때문에 무향 그래프로 이해했고, 도로를 기준으로 최소 비용의 도로들을 선택해서 모든 도시를 이을 수 있어야 하기 때문에, 사이클이 없어야 하는 크루스칼 알고리즘으로 해결했습니다.

 

크루스칼 알고리즘을 공부하기에 앞서 서로소 집합을 트리로 표현하는 부분에 대해 먼저 알고 있어야 합니다. 이와 관련된 내용은 아래 포스트를 참고해주세요 :)

[JAVA] 서로소 집합(Disjoint Sets)과 연산

 

아래는 크루스칼 알고리즘을 이용하여 푼 코드입니다. 시간복잡도를 줄이기 위해 PriorityQueue 를 이용하여 가장 적은 비용의 간선을 추출하고 Union-Find 를 사용하여 사이클을 검사했습니다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {
    private static int N;
    private static int max = 50009;
    private static PriorityQueue<Edge> pq;
    private static int[] parents = new int[max];
    private static int res = 0;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();

        for (int test_case = 1; test_case <= T; test_case++) {
            N = Integer.parseInt(br.readLine()); // 도시 수
            int M = Integer.parseInt(br.readLine()); // 도로 후보 수 (간선 수)
            pq = new PriorityQueue<>();
            res = 0;

            StringTokenizer st;
            for (int i=0; i<M; i++) {
                st = new StringTokenizer(br.readLine());
                int s = Integer.parseInt(st.nextToken());
                int e = Integer.parseInt(st.nextToken());
                int c = Integer.parseInt(st.nextToken());

                // 우선순위큐로 가장 비용 낮은 간선부터
                pq.offer(new Edge(s, e, c)); // s 에서 e 로 갈 때 최소 비용
            }
            makeSet(); // 자기 자신 집합 먼저 만들고

            while (!pq.isEmpty()) {
                Edge e = pq.poll();
                int x = find(e.start); // 대표 원소 찾기
                int y = find(e.end);
                if (x == y) continue; // 사이클이므로 건너뛴다 (사이클 없어야 함)
                else {
                    res += e.cost;
                    union(x, y); // 합집합으로
                }
            }
            sb.append("#").append(test_case).append(" ").append(res).append("\n");
        }
        System.out.println(sb);
    }

    private static void union(int x, int y) {
        if (x > y) {
            parents[x] = parents[y];
        } else {
            parents[y] = parents[x];
        }
    }

    private static int find(int x) {
        if (x == parents[x]) {
            // 자기 자신이면 대표 원소이므로
            return x;
        } else {
            return x = find(parents[x]);
        }
    }

    private static void makeSet() {
        for (int i=1; i<=N; i++) {
            parents[i] = i;
        }
    }

    private static class Edge implements Comparable<Edge> {
        int start;
        int end;
        int cost;

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

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