[Algorithm] 팀 결성

문제

학교에서 학생들에게 0번부터 N번까지 번호를 부여했다. 처음에는 모든 학생이 서로 다른 팀으로 구분되어 총 N+1개의 팀이 존재한다. 이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.

 

1. 팀 합치기 연산은 두 팀을 합치는 연산이다.

2. 같은 팀 여부 확인 연산은 특정한 두 학생이 같은 팀에 속하는지를 확인하는 연산이다.

 

선생님이 M개의 연산을 수행할 수 있을 때, 같은 팀 여부 확인 연산에 대한 연산 결과를 출력하는 프로그램을 작성하시오.

 

입력조건

첫째줄에 N, M이 주어진다.

둘째줄부터 M+1번째 줄까지 각각의 연산이 주어진다.

팀 합치기 연산은 0 a b 형태로 주어진다.

같은 팀 여부 확인 연산은 1 a b 형태로 주어진다.

 

출력 조건

같은 팀 여부 확인 연산에 대해 한줄에 하나씩 YES 혹은 NO를 출력한다.

 

풀이

팀을 합치고, 팀인지 확인하는 서로소 집합 자료구조를 활용한 문제다.

집합을 합치는 연산과 집합을 확인하는 연산을 구현한다.

import java.util.*;

class Main {

  public static int n, m;
  // 팀 번호
  public static int[] team = new int[100001];
  // 결과 저장
  public static List<String> result = new ArrayList<>();

  // 팀 합치기
  public static void unionTeam (int a, int b) {
    a = findTeam(a);
    b = findTeam(b);
    if (a < b) team[b] = a;
    else team[a] = b;
  }

  // 같은 팀 여부 확인
  public static int findTeam (int x) {
    if (team[x] == x) return x;
    return team[x] = findTeam(team[x]);
  }

  public static void main(String[] args) {
    
    Scanner sc = new Scanner(System.in);
    n = sc.nextInt();
    m = sc.nextInt();

    for (int i = 0; i <= n; i++) {
      team[i] = i;
    }

    for (int i = 0; i < m; i++) {
      int f = sc.nextInt();
      int a = sc.nextInt();
      int b = sc.nextInt();
      if (f == 0) {
        // 합치기
        unionTeam(a, b);
      } else {
        // 확인
        result.add((findTeam(a) == findTeam(b)) ? "YES" : "NO");
      }
    }

    for (String r : result) {
      System.out.print(r + " ");
    }
    

  } // end main
}