113. 팀 결성

• 학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다.

  • 처음에는 모든 학생이 서로 다른 팀으로 구분되어 총 N + 1개의 팀이 존재한다.

  • 이 때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.

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

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

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

• 입력조건

  • 첫째 줄에 N, M이 주어진다. M은 입력으로 주어지는 연산의 개수이다. (1 ≤ N, M ≤ 100,000)

  • 다음 M개의 줄에는 각각의 연산이 주어진다.

  • '팀 합치기' 연산은 0 a b 형태로 주어진다. 이는 a번 학생이 속한 팀과 b번 학생이 속한 팀을 합친다는 의미이다.

  • '같은 팀 여부 확인' 연산은 1 a b 형태로 주어진다. 이는 a번 학생과 b번 학생이 같은 팀에 속해 있는지를 확인하는 연산이다.

  • a와 b는 N이하의 양의 정수이다.

• 출력조건

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

1. 서로소 집합 알고리즘을 이용한 풀이

def find_parent(parent, x):
  if parent[x] != x:
    parent[x] = find_parent(parent, parent[x])
    
  return parent[x] #경로 압축 방식

def union_parent(parent, a, b):
  a = find_parent(parent, a)
  b = find_parent(parent, b)

  if a < b: 
    parent[b] = a # 큰수 -> 작은수(부모)
  else:
    parent[a] = b

n, m = map(int, input().split())
parent = [0] * (n + 1)

for i in range(0, n + 1):
  parent[i] = i

result = []

for i in range(m):
  oper, a, b = map(int, input().split())
  #합집합 연산인 경우
  if oper == 0:
    union_parent(parent, a, b)
  elif oper == 1:
    if find_parent(parent, a) == find_parent(parent, b):
      result.append('YES')
    else:
      result.append('NO')

for r in result:
  print(r)

C++ 코드

#include <bits/stdc++.h>

using namespace std;

// 노드의 개수(N)와 연산의 개수(M)
int n, m;
int parent[100001]; // 부모 테이블 초기화

// 특정 원소가 속한 집합을 찾기
int findParent(int x) {
    // 루트 노드가 아니라면, 루트 노드를 찾을 때까지 재귀적으로 호출
    if (x == parent[x]) return x;
    return parent[x] = findParent(parent[x]);
}

// 두 원소가 속한 집합을 합치기
void unionParent(int a, int b) {
    a = findParent(a);
    b = findParent(b);
    if (a < b) parent[b] = a;
    else parent[a] = b;
}

int main(void) {
    cin >> n >> m;

    // 부모 테이블상에서, 부모를 자기 자신으로 초기화
    for (int i = 1; i <= n; i++) {
        parent[i] = i;
    }

    // 각 연산을 하나씩 확인 
    for (int i = 0; i < m; i++) {
        int oper, a, b;
        cin >> oper >> a >> b;
        // 합집합(Union) 연산인 경우
        if (oper == 0) {
            unionParent(a, b);
        }
        // 찾기(Find) 연산인 경우
        else if (oper == 1) {
            if (findParent(a) == findParent(b)) {
                cout << "YES" << '\n';
            }
            else {
                cout << "NO" << '\n';
            }
        }
    }
}

Java 코드

import java.util.*;

public class Main {

    // 노드의 개수(N)와 연산의 개수(M)
    // 노드의 개수는 최대 100,000개라고 가정
    public static int n, m;
    public static int[] parent = new int[100001]; // 부모 테이블 초기화

    // 특정 원소가 속한 집합을 찾기
    public static int findParent(int x) {
        // 루트 노드가 아니라면, 루트 노드를 찾을 때까지 재귀적으로 호출
        if (x == parent[x]) return x;
        return parent[x] = findParent(parent[x]);
    }

    // 두 원소가 속한 집합을 합치기
    public static void unionParent(int a, int b) {
        a = findParent(a);
        b = findParent(b);
        if (a < b) parent[b] = a;
        else parent[a] = b;
    }

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

        n = sc.nextInt();
        m = sc.nextInt();

        // 부모 테이블상에서, 부모를 자기 자신으로 초기화
        for (int i = 1; i <= n; i++) {
            parent[i] = i;
        }

        // 각 연산을 하나씩 확인 
        for (int i = 0; i < m; i++) {
            int oper = sc.nextInt();
            int a = sc.nextInt();
            int b = sc.nextInt();
            // 합집합(Union) 연산인 경우
            if (oper == 0) {
                unionParent(a, b);
            }
            // 찾기(Find) 연산인 경우
            else if (oper == 1) {
                if (findParent(a) == findParent(b)) {
                    System.out.println("YES");
                }
                else {
                    System.out.println("NO");
                }
            }
        }
    }
}