백준 15809 - 전국시대

문제

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

리뷰

유니온 파인드를 이용하여 풀이할 수 있는 간단한 문제였다.

nums는 각 국가별 병력을 저장한다. 유니온 파인드 연산을 활용하여
두 국가가 연합할 경우 연산을 수행하며 루트가 되는 노드의 nums 값에
병력을 합하여 저장해주고, 전쟁이 발발할 경우 승리한 국가의 nums 값을
갱신해주는 방식으로 로직을 구성하였다.

두 나라가 병력을 같을 떄 전쟁을 하면 두 나라 모두 망한 것으로 처리를
해주어야 하기 때문에 이때는 nums 값을 두 나라 다 0으로 처리해주어 구별했다.

마지막으로 루트이며(parent의 값이 음수) 망하지 않은(nums의
값이 0이 아닌) 국가들의 nums 값을 List 에 저장하고 정렬하여 출력하여
답을 구했다.

로직의 시간복잡도는 연합, 전쟁을 처리하는 부분에서 $O(Ma(N))$으로 수렴하며
무난히 제한 조건 1초를 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

import static java.lang.Integer.*;

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

    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());

        nums = new int[N + 1];
        parent = new int[N + 1];

        Arrays.fill(parent, -1);

        int cost;
        for (int i = 1; i <= N; i++) {
            cost = parseInt(br.readLine());
            nums[i] = cost;
        }

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

            q = parseInt(st.nextToken());

            u = parseInt(st.nextToken());
            v = parseInt(st.nextToken());

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

            if (q == 1) {
                if (r1 == r2) continue;

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

            } else {
                if (nums[r1] == nums[r2]) {
                    nums[r1] = nums[r2] = 0;
                } else if (nums[r1] > nums[r2]) {
                    nums[r1] -= nums[r2];
                    nums[r2] = 0;

                    parent[r1]+=parent[r2];
                    parent[r2]=r1;

                } else {
                    nums[r2] -= nums[r1];
                    nums[r1] = 0;

                    parent[r2]+=parent[r1];
                    parent[r1]=r2;
                }
            }
        }

        List<Integer> resultNums = new ArrayList<>();
        for (int i = 1; i < parent.length; i++) {
            if (parent[i] < 0 && nums[i] > 0)
                resultNums.add(nums[i]);
        }

        Collections.sort(resultNums);
        System.out.println(resultNums.size());
        for (Integer resultNum : resultNums) {
            System.out.print(resultNum + " ");
        }
        br.close();
    }

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

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

결과