문제
0부터 N까지의 정수가 각각이 집합 하나를 이루고 있을때 주어진 명령어에 따라 집합을 합친다. 그 사이사이에 정수 A,B가 같은 집합에 속해있는지 묻는다.
풀이
find-union 알고리즘을 새로 배웠다.
parent-child 관계를 통해 같은 root를 갖으면 같은 집합이라고 생각하는 것이다. 코드를 봐보자.
from collections import defaultdict
n,m = map(int,input().split())
parent = {}
for i in range(n+1):
parent[i] = i
instructions = []
for _ in range(m):
instructions.append(list(map(int,input().split())))
def find(a):
if a == parent[a]:
return a
parent[a] = find(parent[a])
return parent[a]
def union(a,b):
a = find(a)
b = find(b)
parent[b] = a
for instruction in instructions:
k,a,b = instruction
if k == 0:
union(a,b)
else:
a = find(a)
b = find(b)
if a == b:
print('YES')
else:
print('NO')
find() 함수의 특징은 root를 return하는 것도 있지만 호출된 노드에서 root까지 가는 길에 있는 모든 node를 root노드의 자식위치까지 올리는 것이다.