[알고리즘] 백준 1717번 - 집합의 표현(Python)

https://www.acmicpc.net/problem/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 를 출력해도 된다)

 

* 예제 입력1

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

 

* 예제 출력1

NO
NO
YES

 

* 풀이

기본적인 Union-Find(유니온-파인드) 문제이다. 0을 입력하였을 경우에는 union연산(합연산)을, 1을 입력하였을 경우에는 find연산을 진행하면 된다. 필자는 union연산을 진행하는 함수, find연산을 진행하는 함수를 따로 만들어 풀었다. find연산을 진행하였을 때, 경로 압축을 이용하면 훨씬 더 소요시간을 절약할 수 있다. 팀 여부를 확인하였을 때(1 입력), 같은 팀이면(parent가 같은 경우) YES를, 다른 팀이면(parent가 다른 경우) NO를 출력해주면 된다. 여기서 주의해야 할 점이 있다. 코드의 재귀의 깊이가 Python이 정한 최대 재귀 깊이보다 초과하는 경우 RuntimeError(RecursionError)가 발생하기 때문에 코드의 재귀 깊이를 한정해야 한다. 재귀 깊이를 한정하는 코드는 다음과 같다.

import sys
sys.setrecursionlimit(10 ** 5)

 다음은 전체 풀이 코드이다.

 

# 집합의 표현
import sys
sys.setrecursionlimit(10 ** 5)

def find_parent(parent, x) :
    if parent[x] != x :
        parent[x] = find_parent(parent, parent[x]) # 경로 압축
        
    return parent[x]

def union_parent(parent, a, b) :
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b :
        parent[b] = a
        
    else :
        parent[a] = b
        
n, m = map(int, input().split())
parent = [0] * (n + 1)
edges = []

for i in range(0, n + 1) :
    parent[i] = i
    
for i in range(m) :
    x, y, z = map(int, input().split())
    edges.append((x, y, z))
    
for edge in edges :
    oper, a, b = edge
    
    if oper == 0 :
        union_parent(parent, a, b)
    
    elif oper == 1 :
        if find_parent(parent, a) == find_parent(parent, b) :
            print("YES")
            
        else :
            print("NO")