백준 14621 - 나만 안되는 연애

문제

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

리뷰

크루스칼을 이용하여 풀이할 수 있는 간단한 문제였다.

주어진 문제 조건에서 모든 대학교로 이동이 가능(2번 조건)과 최단 거리(3번 조건)은
MST를 구해야 한다는 말로 해석할 수 있다. 다만, 1번 조건에 따라 두 정점이 같은
유형의 대학교(남초==남초, 여초==여초)가 아닌지 확인하는 부분이 크루스칼을
로직에 추가적으로 필요했다.

추가 조건을 구현하기 위해 boolean 타입의 gender 배열을 정의하여
true일 때를 남초, false 일 때를 여초로 설정하고 크루스칼 로직에서 두 정점의
대학교 유형이 다를 때만 해당 간선을 채택할 것인지 말 것인지 판단하게 하여 답을
구했다.

로직의 시간복잡도는 크루스칼의 $O(NlogM)$으로 수렴하고 이는 $N=1,000$,
$M=10,000$인 최악의 경우에도 무난히 제한 조건 2초를 통과한다.

코드

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

import static java.lang.Integer.*;

public class Main {
    static int[] parent;
    static boolean[] gender;
    static PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingInt(e -> e.w));

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = parseInt(st.nextToken());
        int M = parseInt(st.nextToken());

        parent = new int[N + 1];
        gender = new boolean[N + 1];

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i < gender.length; i++) {
            gender[i] = st.nextToken().equals("M");
        }

        int u, v, w;
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());
            w = parseInt(st.nextToken());

            pq.offer(new Edge(u, v, w));
        }

        System.out.println(kruskal(N));
        br.close();
    }

    static int kruskal(int N) {
        Arrays.fill(parent, -1);
        int r1, r2, selected = 0, total = 0;

        while (!pq.isEmpty() && selected < N - 1) {
            Edge e = pq.poll();

            if (gender[e.u] == gender[e.v]) continue;

            r1 = find(e.u);
            r2 = find(e.v);

            if (r1 == r2) continue;

            if (parent[r1] < parent[r2]) {
                parent[r1] += parent[r2];
                parent[r2] = r1;
            } else {
                parent[r2] += parent[r1];
                parent[r1] = r2;
            }

            selected++;
            total += e.w;
        }

        return selected == N - 1 ? total : -1;
    }

    static int find(int u) {
        if (parent[u] < 0) return u;

        return parent[u] = find(parent[u]);
    }

    static class Edge {
        int u, v, w;

        public Edge(int u, int v, int w) {
            this.u = u;
            this.v = v;
            this.w = w;
        }
    }
}

결과