백준 1717 [집합의 표현] 파이썬

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

# 백준 1717 집합의 표현 (Union-Find 알고리즘 활용)

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)  # 재귀 깊이 제한을 넉넉히 설정

# n: 원소 개수, m: 연산 횟수
n, m = map(int, input().split())

# 각 원소의 부모를 저장하는 배열
# 처음에는 자기 자신을 부모로 가짐
parent = [i for i in range(n + 1)]


# find 연산 (경로 압축 적용)
# 원소 x의 루트(parent[x] == x인 노드)를 찾는 함수
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])  # 경로 압축: 루트를 바로 부모로 연결
    return parent[x]


# union 연산
# 두 원소가 속한 집합을 합치는 함수
def union(a, b):
    a = find(a)  # a의 루트
    b = find(b)  # b의 루트

    if a != b:   # 루트가 다르면 연결
        parent[a] = b


# 두 원소가 같은 집합에 속하는지 확인하는 함수
def is_in(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        return False
    else:
        return True


# m번의 연산 처리
for _ in range(m):
    oper, A, B = map(int, input().split())

    if oper == 0:
        # 합집합(union) 연산
        union(A, B)
    else:
        # 같은 집합에 속하는지 확인
        if is_in(A, B):
            print('YES')
        else:
            print('NO')

UNION-FIND가 익숙해진다.