[🥇4 / 백준 1717 / 파이썬] 집합의 표현

 

1717번: 집합의 표현

 

문제

초기에 {0}, {1}, {2}, ... {n} 이 각각 n+1개의 집합을 이루고 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

입력

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

 

예제 입력  7 8 0 1 3 1 1 7 0 7 6 1 7 1 0 3 7 0 4 2 0 1 1 1 1 1

예제 출력  NO NO YES

 

풀이

 

분리 집합(Disjoint set)에 대해 먼저 알아보자.

분리 집합이란 교집합이 존재하지 않는 둘 이상의 집합을 의미한다.

 

보통 트리 구조를 이용한 유니온 파인드를 이용해 분리 집합을 구현한다.

여기서 Union은 두 분리 집합을 하나로 합치는 것, Find는 어떤 자식 노드의 최정상 부모 노드를 찾아 해당 원소가 속한 집합을 찾는 연산을 의미한다.

 

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**8)

def union(x,y):
    x = find(x)
    y = find(y)
    # 번호가 낮은 순서대로 트리의 상위 부모 노드로 위치하도록
    if x < y:
        parents[y] = x
    else:
        parents[x] = y

def find(k):
    if k != parents[k]:
        parents[k] = find(parents[k])
    return parents[k]

n,m = map(int,input().split())
parents = [i for i in range(n+1)]
for _ in range(m):
    c,a,b = map(int,input().split())
    if c:
        print("YES" if find(a) == find(b) else "NO")
    else:
        union(a,b)