[이것이 코딩테스트다] 팀 결성

🗃️문제 설명

학교에서 학생들에게 0번부터 N번까지의 번호를 부여했다.
처음에는 모든 학생이 서로 다른 팀으로 구분되어, 총 N + 1개의 팀이 존재한다.
이때 선생님은 '팀 합치기' 연산과 '같은 팀 여부 확인' 연산을 사용할 수 있다.
'팀 합치기' 연산은 두 팀을 합치는 연산이다.
‘같은 팀 여부 확인' 연산은 특정한 두 학생이 같은 팀에 속하는지를 확인하는 연산이다.
선생님이 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로 결과를 출력한다.

🖥️코드

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
parent = [i for i in range(N+1)] # 부모 테이블

def find(x):
    # parent[x] : 현재 원소 x의 부모 노드
    # x : 현재 원소
    # 즉, 현재 원소가 루트가 아님
    if parent[x] != x:
        parent[x] = find(parent[x]) # 현재 원소의 부모를 계속 따라 올라간다.
    return parent[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a < b:   # 작은 쪽으로 편성
        parent[b] = a
    else:
        parent[a] = b

for _ in range(M):
    type, a, b = map(int, input().split())
    if type == 0:
        union(a, b)
    else:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")

🔒문제 출처 및 참고 코드

이것이 코딩테스트다 with 파이썬 - 팀 결성
모범 답안 - 팀 결성