[SWEA] 3289 서로소집합 - JAVA

1. 문제

SW Expert Academy

초기에 {1}, {2}, ... {n} 이 각각 n개의 집합을 이루고 있다.
여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다.
두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다.
a와 b는 n 이하의 자연수이며 같을 수도 있다.

---

2. 풀이

import java.io.*;
import java .util.*;

public class D4_3289_서로소집합 {
   static int[] sets, size;
   public static void main(String[] args) throws IOException {
      BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
      int T = Integer.parseInt(in.readLine());

      StringBuilder sb = new StringBuilder();

      for(int tc=1;tc<=T;tc++){
         sb.append("#").append(tc).append(" ");

         StringTokenizer st = new StringTokenizer(in.readLine());
         int N = Integer.parseInt(st.nextToken());
         int M = Integer.parseInt(st.nextToken());

         size = new int[N+1];
         sets = new int[N+1];
         for(int i=1;i<=N;i++) sets[i] = i;


         for(int m=0;m<M;m++){

            st = new StringTokenizer(in.readLine());

            int operator = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            switch(operator){
               case 0://합집합
                  union(a,b);
                  break;
               case 1://같은집합인지 확인
                  if(find(a)==find(b)) sb.append(1);
                  else sb.append(0);
                  break;
            }
         }
         sb.append("\n");
      }
      System.out.print(sb);
   }

   private static int find(int x){
      if(sets[x]==x) return x;
      sets[x] = find(sets[x]);
      return sets[x];
   }

   private static void union(int a, int b){
      int pa = find(a);
      int pb = find(b);

      if(pa==pb) return;

      if(size[pa]>size[pb]) sets[pb] = pa;
      else if(size[pa]==size[pb]) size[pb]++;

      if(size[pb]>size[pa]) sets[pa] = pb;

   }
}

---

3. 결과