[SWEA] 1251. 하나로

문제 링크

[SWEA] 1251. 하나로

풀이 과정

크루스칼 알고리즘이 적용된 기본 문제입니다.

크루스칼 알고리즘에 대한 설명은 링크에 있습니다.

최소 신장 트리를 구현하는 문제입니다. 일반적인 MST 문제와는 다르게 시작 정점, 도착 정점, 비용이 주어지는 것이 아니라 모든 정점의 좌표가 주어지고, 이를 토대로 비용을 계산해내는 것이 이 문제의 특징입니다.

간선을 최대한으로 갖는 완전 그래프의 간선 최대 간선 수는 N * (N-1) / 2 입니다. 이를 활용해 모든 간선의 비용을 계산해 배열에 저장해 두었습니다.

코드

import java.util.Arrays;
import java.util.Comparator;
import java.util.Scanner;

public class Solution {
    static int[] parent;
    static int[] rank;
    static int N;
    static int[] X, Y;
    static double E;
    static Edge[] edges;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int TC = sc.nextInt();
        for (int tc = 1; tc <= TC; tc++) {
            N = sc.nextInt();

            parent = new int[N];
            rank = new int[N];

            edges = new Edge[N * (N - 1) / 2];

            X = new int[N];
            Y = new int[N];

            for (int i = 0; i < N; i++) {
                X[i] = sc.nextInt();
            }

            for (int i = 0; i < N; i++) {
                Y[i] = sc.nextInt();
            }

            E = sc.nextDouble();
            int cnt = 0;
            for (int i = 0; i < N - 1; i++) {
                for (int j = i + 1; j < N; j++) {
                    if (cnt == N * (N - 1) / 2 + 1)
                        break;
                    edges[cnt++] = new Edge(i, j, dist(i, j));
                }
            }

            for (int i = 0; i < N; i++)
                makeSet(i);

            Arrays.sort(edges, new Comparator<Edge>() {
                @Override
                public int compare(Edge o1, Edge o2) {
                    return Double.compare(o1.cost, o2.cost);
                }
            });

            cnt = 0;
            double answer = 0;
            for (Edge e : edges) {
                int a = e.start;
                int b = e.dest;

                if (find(a) == find(b)) continue;

                union(a, b);
                answer += e.cost;
                cnt++;
            }

            System.out.println("#" + tc + " " + Math.round(answer));
        }
    }

    static double dist(int i, int j) {
        double d = Math.sqrt(Math.pow(X[i] - X[j], 2) + Math.pow(Y[i] - Y[j], 2));
        return E * d * d;
    }

    static void makeSet(int x) {
        parent[x] = x;
    }

    static int find(int x) {
        if (x == parent[x]) return x;
        else return parent[x] = find(parent[x]);
    }

    static void union(int x, int y) {
        int px = find(x);
        int py = find(y);

        if (rank[px] > rank[py])
            parent[py] = px;
        else {
            parent[px] = py;

            if (rank[px] == rank[py])
                rank[px]++;
        }
    }

    static class Edge {
        int start;
        int dest;
        double cost;

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

        @Override
        public String toString() {
            return "Edge{" +
                    "start=" + start +
                    ", dest=" + dest +
                    ", cost=" + cost +
                    '}';
        }
    }
}