[snippet] union-find.py

백준 1717번. 집합의 표현에 대한 풀이를 Union-Find의 스니펫 느낌으로 작성.

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

n, m = map(int, sys.stdin.readline().rstrip().split())
parent = [i for i in range(n+1)]

#########################################

def get_parent(x):
    if parent[x] == x:
        return x
    parent[x] = get_parent(parent[x]) # get_parent 거슬러 올라가면서 parent[x] 값도 갱신
    return parent[x]

def union_parent(a, b):
    a = get_parent(a)
    b = get_parent(b)
    
    if a < b: # 작은 쪽을 parent로 합의하자
        parent[b] = a
    else:
        parent[a] = b        

def same_parent(a, b):
    return get_parent(a) == get_parent(b)

#########################################

for _ in range(m):
    op, a, b = map(int, sys.stdin.readline().rstrip().split())
    if op == 0:
        union_parent(a, b)
    else:
        print('YES' if same_parent(a, b) else 'NO')

주의사항

same_parent 는 잘 작동하겠지만, parent가 아직 통일된 상태는 아니다.
parent 배열을 직접 이용하려는 경우에는 접근 이전에 다음의 코드를 반드시 실행하자.

for i in range(1, n+1):
	get_parent(i)

틀린 예시