백준 16562 - 친구비

문제

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

리뷰

유니온 파인드로 풀이할 수 있는 문제였다.

문제의 답은 각 친구 무리에서 가장 작은 친구비의 합을 구하는 것이다.
따라서 parent를 $2 \times (N+1)$의 2차원 배열로 정의하고 parent[1] 에는
루트 노드 혹은 자식 노드의 수(음수값)를 저장하며, parent[0]에는 해당
분리 집합에서 가장 작은 친구비를 저장하도록 구성하였다.

유니온 파인드 연산을 수행하며 parent[0],parent[1]의 값을 갱신해준 후
루트에 해당하는(parent[1]의 값이 음수) 경우만 parent[0]의 값을
더해주면 전체 친구 비용을 도출하여 k와 비교해 답을 판별할 수 있다.

로직의 시간복잡도는 유니온 파인드 연산이 수행되는 부분의 $O(Ma(N))$으로
수렴하고 $N=M=10,000$인 최악의 경우에도 무난히 제한 조건 2초를 통과한다.

코드

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

import static java.lang.Integer.*;

public class Main {
    static int[][] parent;
    static int[] cost;

    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());
        int k = parseInt(st.nextToken());

        parent = new int[2][N + 1];
        cost = new int[N + 1];

        Arrays.fill(parent[1], -1);

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i < cost.length; i++) {
            cost[i] = parseInt(st.nextToken());
            parent[0][i] = cost[i];
        }

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

            int r1 = find(u);
            int r2 = find(v);

            if (r1 == r2) continue;

            if (parent[1][r1] < parent[1][r2]) {
                parent[1][r1] += parent[1][r2];
                parent[1][r2] = r1;
                parent[0][r1] = Math.min(parent[0][r1], parent[0][r2]);
            } else {
                parent[1][r2] += parent[1][r1];
                parent[1][r1] = r2;
                parent[0][r2] = Math.min(parent[0][r1], parent[0][r2]);
            }
        }

        int totalCost = 0;

        for (int i = 1; i < parent[1].length; i++)
            if (parent[1][i] < 0)
                totalCost += parent[0][i];

        System.out.println(totalCost > k ? "Oh no" : totalCost);
        br.close();
    }

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

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

결과