1717 집합의 표현 [union-find]

문제 예제가 좀 애매해서 문제를 이해하는데 좀 시간이 걸렸다. 결국 알고리즘분류가 disjoint-set 인걸보고 이해했다.

이 문제는 단순히 set을 이용해서 풀수도 있지만 입력제한에서 볼수있듯이 시간초과가 날수밖에없다.

그렇다면 disjoint-set 다른말로 union-find 자료구조를 이용하면 해결할수 있다.

union-find는 크루스칼 알고리즘에도 사용되는 구조로 즉, 원소가 같은 루트값을 가지는지 효과적으로 체크해주는 구조로 쉽게 생각하면 쉽다

여기서 parent가 각 부모를 나타내는 리스트고 rank는 union을 최적화하기 위해 각 트리의 레벨을 설정함으로써 

union 연산시 트리의 레벨이 낮은 쪽이 큰쪽밑으로 붙도록 하는 것이다. 즉 최적화를 위해 만든 리스트라고 보면된다. (없어도 돌아가긴 한다)

import sys
read = sys.stdin.readline


def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]


def union(a, b):
    root1 = find(a)
    root2 = find(b)
    if root1 != root2:
        if rank[root1] > rank[root2]:
            parent[root2] = root1
        else:
            parent[root1] = root2
            if rank[root1] == rank[root2]:
                rank[root2] += 1

n, m = map(int, read().split())
parent = [i for i in range(n+1)]  # 부모
rank = [0 for i in range(n+1)]    # 트리의 깊이

for i in range(m):
    o, a, b = map(int, read().split())
    if o:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")

    else:
        union(a, b)