[#알고리즘/Union-Find/[백준]1717번: 집합의 표현(python)]

Question

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

Solution

기본적인 Union-Find 알고리즘 동작 과정이다. 처음에 sys.setrecursionlimit() 없이 실행하였더니 런타임 에러가 발생하였다. 재귀를 사용해 문제를 풀 때 런타임 에러가 발생하는 경우에는, 재귀의 최대 깊이를 적절히 설정해줘야 한다. 주의할 점은, PyPy에서는 sys.setrecursionlimit()으로 임의로 재귀의 최대 깊이를 설정할 수 없다.

Code

import sys
sys.setrecursionlimit(10**5)
input = sys.stdin.readline

# 특정 원소가 속한 집합 찾기 (즉, 재귀적으로 부모를 찾으며 원소의 루트 노드(속한 집합) 찾기)
def find_parent(parent, x) :
    if parent[x] != x :
        parent[x] = find_parent(parent, parent[x])
    return parent[x] 
        
# union 연산
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(n+1) :
    parent[i] = i

# union 연산 각각 수행
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) :
            print('YES')
        else :
            print('NO')