[백준]#15809 전국시대

문제

전국시대엔 N개의 국가가 존재한다. 각 국가는 1부터 N까지의 번호를 가지고 있다.

또한, 모든 국가는 각자 자신의 국가의 힘을 상징하는 병력을 가지고 있다. 이때 M개의 기록이 주어진다. 각각의 기록은 다음과 같다.

1. 동맹 - 두 나라가 서로 동맹을 맺는다. 두 나라의 병력이 하나로 합쳐진다.
2. 전쟁 - 두 나라가 서로 전쟁을 벌인다. 병력이 더 많은 나라가 승리하며 패배한 나라는 속국이 된다. 이때 남은 병력은 승리한 나라의 병력에서 패배한 나라의 병력을 뺀 수치가 된다. 두 나라의 병력이 같을 경우 두 나라 모두 멸망한다.

모든 나라는 정직하기 때문에 내 동맹의 동맹도 나의 동맹이고, 내 동맹이 적과 전쟁을 시작하면 같이 참전한다. 속국인 경우도 동맹의 경우와 마찬가지이다.

따라서, 전쟁에서 진 국가와 동맹인 다른 국가 또한 전쟁에서 이긴 국가의 속국이 된다.

모든 기록이 끝났을 때 남아있는 국가의 수를 출력하고, 그 국가들의 남은 병력의 수를 오름차순으로 출력하는 프로그램을 작성하시오.

단, 여러 국가가 서로 동맹이거나 속국 관계인 경우는 한 국가로 취급한다.

입력

첫 번째 줄에 국가의 수를 나타내는 N과 기록의 수 M이 주어진다. (1 ≤ N, M ≤ 100,000)

두 번째 줄 부터 N개의 줄에 걸쳐 i번째 국가의 병력 Ai (1 ≤ i ≤ N)가 자연수로 주어진다. (1 ≤ Ai ≤ 10,000)

다음 M개의 줄에는 기록이 3개의 정수 O, P, Q로 주어진다. O가 1인 경우 P, Q가 서로 동맹을 맺었음을 의미하고, O가 2인 경우 P, Q가 서로 전쟁을 벌였음을 의미한다.

동맹끼리 다시 동맹을 맺거나 전쟁하는 입력은 주어지지 않는다.

출력

첫째 줄에 남아있는 국가의 수를 출력한다.

다음 줄에 각 국가의 남은 병력의 수를 띄어쓰기를 간격으로 오름차순으로 출력한다.

예제 입력 1

5 3
10
20
30
40
50
1 1 2
1 3 4
2 1 3

예제 출력 1

2
40 50

풀이

이 문제는 find-union 알고리즘을 이용해서 풀 수 있었다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashSet;
import java.util.PriorityQueue;

public class Main {
    static int[] parent;
    static int[] num;
    static HashSet<Integer>[] set;

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] input = br.readLine().split(" ");
        int N = Integer.parseInt(input[0]);
        int M = Integer.parseInt(input[1]);
        parent = new int[N+1];
        num = new int[N+1];
        set = new HashSet[N+1];
        for(int i=1; i<=N; i++) {
            set[i] = new HashSet<>();
            set[i].add(i);
        }

        for(int i=1; i<=N; i++)
            parent[i] = i;

        for(int i=1; i<=N; i++)
            num[i] = Integer.parseInt(br.readLine());

        for(int i=0; i<M; i++) {
            input = br.readLine().split(" ");
            int O = Integer.parseInt(input[0]);
            int P = Integer.parseInt(input[1]);
            int Q = Integer.parseInt(input[2]);

            union(O, P, Q);
        }

        PriorityQueue<Integer> pq = new PriorityQueue<>();
        for(int i=1; i<=N; i++) {
            if(parent[i]!=0)
                pq.add(num[i]);
        }

        System.out.println(pq.size());
        while(!pq.isEmpty())
            System.out.print(pq.poll()+" ");
    }

    public static void union(int flag, int a, int b) {
        a = find(a);
        b = find(b);

        if(flag==1) {           //동맹 맺는 경우
            for(int x : set[b]) {
                set[a].add(x);
                num[a] += num[x];
                num[x] = 0;
            }
            set[b].clear();
            parent[b] = a;
        }

        else {
            if(num[a]==num[b]) {    //전쟁시 병력이 같을 때
                num[a] = 0;
                num[b] = 0;
                set[a].clear();
                set[b].clear();
            }

            else if(num[a]>num[b]) {    //전쟁시 병력이 한쪽이 더 클 때
                num[a] -= num[b];
                num[b] = 0;
                for(int x : set[b]) {
                    set[a].add(x);
                }
                set[b].clear();
                parent[b] = a;
            }

            else {
                num[b] -= num[a];
                num[a] = 0;
                for(int x : set[a]) {
                    set[b].add(x);
                }
                set[a].clear();
                parent[a] = b;
            }
        }
    }

    public static int find(int a) {
        if(parent[a]==a)
            return a;

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