[BaekJoon] 14621 나만 안되는 연애 (Java)

1.  문제 링크

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

2.  문제

요약

• 미팅 애플리케이션을 만드려고 하는데, 미팅 앱은 대학생을 타겟으로 만들어졌으며 대학교간의 도로 데이터를 수집하여 만들어집니다.
• 이 앱은 사용자들을 위해 사심 경로를 제공하는데, 이 경로는 세 가지 특징이 있습니다.
  1. 사심 경로는 사용자들의 사심을 만족시키기 위해 남초 대학교와 여초 대학교들을 연결하는 도로로만 이루어져 있습니다.
  2. 사용자들이 다양한 사람과 미팅할 수 있도록 어떤 대학교에서든 모든 대학교로 이동이 가능한 경로입니다.
  3. 시간을 낭비하지 않고 미팅할 수 있도록 이 경로의 길이는 최단 거리가 되어야 합니다.
• 주어지는 거리 데이터를 이용하여 사심 경로의 길이를 구하는 문제입니다.
• 입력: 첫 번째 줄에 2보다 크거나 같고 1,000보다 작거나 같은 학교의 수 N과 1보다 크거나 같고 10,000보다 작거나 같은 도로의 개수 M이 주어지고 두 번째 줄에 각 학교가 남초 대학교라면 M, 여초 대학교라면 W가 주어집니다. 세 번째 줄부터 M개의 줄에 u, v, d가 주어지는데, u 학교와 v 학교가 연결되어 있으며 이 거리는 d임을 나타냅니다.
  • u, v는 1보다 크거나 같고 N보다 작거나 같습니다.
  • d는 1보다 크거나 같고 1,000보다 작거나 같습니다.
• 출력: 첫 번째 줄에 만든 앱의 경로 길이를 출력합니다. 모든 학교를 연결하는 경로가 없을 경우 -1을 출력합니다.

3.  소스코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.StringTokenizer;

public class Main {
    static int N, M;
    static int[][] edges;
    static char[] genders;
    static int[] parents;

    static void input() {
        Reader scanner = new Reader();

        N = scanner.nextInt();
        M = scanner.nextInt();

        edges = new int[M][3];
        genders = new char[N + 1];
        parents = new int[N + 1];
        for(int school = 1; school <= N; school++) parents[school] = school;

        for(int idx = 1; idx <= N; idx++) genders[idx] = scanner.next().charAt(0);

        for(int idx = 0; idx < M; idx++) {
            int school1 = scanner.nextInt(), school2 = scanner.nextInt(), distance = scanner.nextInt();

            edges[idx][0] = school1;
            edges[idx][1] = school2;
            edges[idx][2] = distance;
        }
    }

    static void solution() {
        Arrays.sort(edges, new Comparator<int[]>() {
            @Override
            public int compare(int[] e1, int[] e2) {
                return e1[2] - e2[2];
            }
        });

        ArrayList<int[]> minEdges = kruskal();
        if(minEdges.size() < N - 1) System.out.println(-1);
        else {
            int distance = 0;
            for(int[] edge : minEdges) distance += edge[2];
            System.out.println(distance);
        }
    }

    static ArrayList<int[]> kruskal() {
        ArrayList<int[]> minEdges = new ArrayList<>();

        int index = 0;
        for(int count = 0; count < N - 1; count++) {
            if(index == edges.length) break;

            int[] edge = edges[index];

            if(genders[edge[0]] == genders[edge[1]] || isSameParent(edge[0], edge[1])) {
                count--;
            } else {
                union(edge[0], edge[1]);
                minEdges.add(edge);
            }
            index++;
        }

        return minEdges;
    }

    static int findParent(int school) {
        if(school == parents[school]) return school;
        return parents[school] = findParent(parents[school]);
    }

    static void union(int school1, int school2) {
        int parent1 = findParent(school1), parent2 = findParent(school2);

        if(parent1 != parent2) {
            if(parent1 < parent2) parents[parent2] = parent1;
            else parents[parent1] = parent2;
        }
    }

    static boolean isSameParent(int school1, int school2) {
        int parent1 = findParent(school1), parent2 = findParent(school2);

        return parent1 == parent2;
    }

    public static void main(String[] args) {
        input();
        solution();
    }

    static class Reader {
        BufferedReader br;
        StringTokenizer st;

        public Reader() {
            br = new BufferedReader(new InputStreamReader(System.in));
        }

        String next() {
            while(st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            return st.nextToken();
        }

        int nextInt() {
            return Integer.parseInt(next());
        }
    }
}

4.  접근

• 주어진 도로들에서 크루스칼 알고리즘을 통해 N - 1개의 도로를 선택합니다.
  • 우선 주어진 도로들을 거리에 대해 오름차순으로 정렬합니다.
  • 각 도로들을 순회하며 N - 1개의 도로를 선택합니다.
    • 만약 현재 도로의 두 학교가 둘 다 남초 대학교이거나 여초 대학교일 경우 혹은 이미 두 대학교가 연결되어 있는 경우 다음 도로를 확인합니다.
      • 그렇지 않다면 두 도로를 연결하고 ArrayList에 해당 도로를 넣어줍니다.
      • N - 1개의 도로를 골랐거나 모든 도로를 순회했다면 도로들을 담은 ArrayList를 반환합니다.
  • 만약 반환한 ArrayList에 N - 1개의 도로가 들어있지 않다면 모든 대학교를 연결할 수 없다는 뜻이므로 -1을 출력하고 그렇지 않다면 ArrayList 안에 있는 N - 1개 도로의 거리를 더하여 출력합니다.