[백준 1717] 집합의 표현(Python)

https://blog.naver.com/kks227/220791837179 에서 서로소 집합을 공부하고, 여기에서 추천한 문제를 풀었다.

1. 요약

초기에 {0}, {1}, {2}, .., {n}가 있다.(i.e. n+1개의 집합) 여기서 합집합 연산과, 두 원소가 같은 집합에 포함되어있는지 확인하고 출력하는 연산을 수행하는 프로그램을 만드는 문제이다.

2. 코드

Python의 경우, 재귀 깊이를 설정을 해야한다.

import sys

n, m = map(int, sys.stdin.readline().rstrip().split())
p = [-1 for i in range(n+1)]
sys.setrecursionlimit(10**5)

def find(node):
    global p
    if p[node] < 0: return node
    p[node] = find(p[node])
    return p[node]

def merge(a, b):
    global p
    pa = find(a)
    pb = find(b)

    if pa == pb: return
    p[pa] = pb
    

for i in range(m):
    cmd, a, b = map(int, sys.stdin.readline().rstrip().split())

    if cmd == 0:
        merge(a, b)

    else:
        if find(a) == find(b): print('YES')
        else: print('NO')