[백준] 집합의 표현

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

1717번: 집합의 표현

 

union-find 알고리즘을 사용하여 분리 집합을 표현하였다.

예를 들어, {1, 2, 3}, {4, 5} 인 집합에서 '0 3 5' 의 입력을 받으면 root[1] = 5, root[2] = 5, root[3] = 5, root[4] = 5, root[5] = 5로 하여 root[i] = 5가 나오는 i는 같은 집합의 원소라고 할 수 있다. 5가 아닌 3이라고 정해도 상관없다.

 

import sys

sys.setrecursionlimit(10**6)

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

def union(a, b):
    a = find(a)
    b = find(b)
    root[a] = b

def find(a):
    if root[a] == a:
        return a
    else:
        root[a] = find(root[a])
        return root[a]

for i in range(m):
    o, a, b = map(int, sys.stdin.readline().split())
    
    if o == 0:
        union(a, b)
    else:
        if find(a) == find(b):
            print('YES')
        else:
            print('NO')