<분리집합> 1717번 집합의 표현 with 파이썬

 

문제

 초기에 {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 를 출력해도 된다)

정답비율

28.739%

알고리즘 분류

 #자료구조 #분리집합 #서로소집합 #unionfind #집합

 

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

def find_parent(parent, k):
    if k != parent[k]:
        parent[k] = find_parent(parent, parent[k])
    return parent[k]

def union_find(parent, x, y):
    x = find_parent(parent, x)
    y = find_parent(parent, y)
    if x < y :
        parent[y] = x
    else :
        parent[x] = y

n, m = map(int, input().split())
parent = [i for i in range(n+1)]
for _ in range(m):
    z, a, b = map(int, input().split())
    if z == 0 :
        union_find(parent, a, b)
    else :
        if find_parent(parent, a) == find_parent(parent, b):
            print('YES')
        else :
            print('NO')

 

 - union-find 자료구조 문제이다. 둘 중 루트노드가 작은 수를 공통으로 갖도록 합치는 연산인 union_find와 루트노드를 찾아가는 find_parent를 먼저 정의했다.

 

 - 그 이후로는 간단했다. 첫 번째 인자가 0이면 union_find를 수행하고, 1이면 find_parent를 이용해서 비교했다.

 

 

---

 

참고

 

1717번: 집합의 표현