백준 15789 - CTP 왕국은 한솔 왕국을 이길 수 있을까?

문제

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

리뷰

유니온 파인드로 풀이할 수 있는 간단한 문제였다.
자세한 로직에 관한 설명은 코드에 주석으로 첨부하였다.

로직의 시간복잡도는 복잡도가 가장 큰 연결 정보를 처리하는 부분에서
$O(Ma(N))$으로 수렴하며 $N=100,000$,$M=200,000$인 최악의 경우에도
무난히 제한 조건 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[] 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());

        parent=new int[N];
        Arrays.fill(parent, -1);

        int u, v;

        /**
         * 주어지는 정점 번호의 범위가 1~N까지인데
         * 0~N-1로 사용하기 위해 모든 번호 입력을 받을 때
         * 1을 빼주었다.
         */
        
        while(M-->0){
            st=new StringTokenizer(br.readLine());
            u=parseInt(st.nextToken())-1;
            v=parseInt(st.nextToken())-1;

            union(u, v);
        }

        int C, H, K;
        st=new StringTokenizer(br.readLine());
        C=parseInt(st.nextToken())-1;
        H=parseInt(st.nextToken())-1;
        K=parseInt(st.nextToken());

        int hRoot=find(H);

        /**
         * 한솔 왕국과 루트가 같지 않은 정점만 선별하여
         * list에 담는다.
         */
        
        int root;
        List<Pair> list=new ArrayList<>();
        for(int i=0; i<parent.length; i++){ // O(N)
            root=find(i);

            if(root==hRoot) continue;

            list.add(new Pair(i, parent[i]));
        }

        /**
         * 최적화된 유니온 파인드 로직을 사용하여 parent 배열의 값이
         * 더 작을 수록 많은 노드들을 자식으로 하는 루트이다.
         * 자식 노드의 수가 많은 루트가 앞으로 오게끔 val 기준으로
         * 리스트를 정렬하였다.
         */
        list.sort(Comparator.comparingInt(p -> p.val));

        /**
         * CTP 왕국과 같은 집합에 속하지 않으며 더 많은 자식 노드를
         * 가지는 노드를 우선으로 최대 K번 union 연산을 수행한다.
         */
        int r1, r2;
        for (Pair pair : list) {
            if(K==0) break;

            r1=find(C);
            r2=find(pair.idx);

            if(r1==r2) continue;

            union(r1, r2);
            K--;
        }

        System.out.println(Math.abs(parent[find(C)]));
        br.close();
    }

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

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

    static void union(int u, int v){
        int r1=find(u);
        int r2=find(v);

        if(r1==r2) return;

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

    static class Pair{
        int idx, val;

        public Pair(int idx, int val) {
            this.idx = idx;
            this.val = val;
        }
    }
}

결과