카카오 예선문제 숏코딩을 풀다가 도저히 어떻게 접근하는지 감이 안잡혀서 Union-find로 풀어야 한다는 것을 알게되어 대표적인 문제로 연습한 문제이다.
문제 이해 및 적용하기
기본적인 문제 답게 따로 설정한 스토리는 없고 주어진 n, m 에서 n개의 집합을 m번의 연산을 주고 그 연산동안 union시키거나 입력받은 a, b가 같은 집합안에 있는지 출력하는 형식이다.
parent = [i for i in range(n+1)]
def find(c):
if(parent[c] == c): return c
parent[c] = find(parent[c])
return parent[c]
def union(ca,cb):
rootOfca = find(ca)
rootOfCb = find(cb)
if not rootOfca == rootOfCb:
parent[rootOfCb] = rootOfca
def check(ca,cb):
if find(ca) == find(cb):
return True
else:
return False
disjoint set 을 union find해주기 위한 코드를 위와 같이 find, union, check를 작성해주었다.
여기서 주의할 점이 find를 parent[c] = find(parent[c]) 위와 같이 최적화 과정을 거쳐 주어야 python3 시간초과가 뜨지 않는다.
본인은 처음에
while c != parent[c]: c = parent[c] 이 코드를 사용하여서 시간초과가 나게 되었다.
주의할 점
- 경로 최신화를 통한 최적화를 해주어야 python3에서 시간초과가 뜨지 않는다.
- 경로 최신화를 하지 않았을 때는 pypy3에서는 통과되지만 python3에서는 통과되지 않는다
전체코드
import sys
input = sys.stdin.readline
# disjoint set & union-find
n, m = map(int,input().split())
parent = [i for i in range(n+1)]
def find(c):
if(parent[c] == c): return c
parent[c] = find(parent[c])
return parent[c]
def union(ca,cb):
rootOfca = find(ca)
rootOfCb = find(cb)
if not rootOfca == rootOfCb:
parent[rootOfCb] = rootOfca
def check(ca,cb):
if find(ca) == find(cb):
return True
else:
return False
# finding answer
ans = []
for i in range(m):
flag, a, b = map(int,input().split())
if flag == 0 :
union(a,b)
else:
if check(a,b):
ans.append("YES")
else:
ans.append("NO")
for i in range(len(ans)):
print(ans[i])