[백준 1717번] 집합의 표현 - 파이썬(python)

1717번: 집합의 표현

분류

• 자료 구조
• 분리 집합

문제

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

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

입력

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

출력

1로 시작하는 입력에 대해서
a와 b가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

풀이

• '집합을 합친다.'와 '두 원소가 같은 집합에 있는지 확인한다'를 보면 유니온 파인드를 사용하는 문제라는 것을 알 수 있다.
• 따라서 부모를 찾는 find 함수와 집합을 합치는 union 함수를 만들어 풀이한다.
• k가 0일 경우는 union을 1일 경우는 find를 한다.
• k가 1일 경우는 입력된 a와 b의 부모를 찾아 같다면 "YES"를 다르다면 "NO"를 출력한다.
• 이 풀이는 파이썬 재귀 제한에 걸리기 때문에 sys.setrecursionlimit(1000000)를 사용하여 재귀 횟수르르 늘려준다.
• 또한 부모를 매번 찾을 경우 시간초과가 발생하기 떄문에 find할 때 미리 부모노드를 저장한다.

코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)

# 부모 찾기
def find(x):
    if parent[x] != x:
        # return find(parent[x])
        # 시간을 줄이기 위해 부모를 찾음과 동시에 저장해야 한다.
        parent[x] = find(parent[x])
    return parent[x]

# 합집합 만들기
def union(a,b):
    a = find(a)
    b = find(b)
    if a<b:
        parent[b]=a
    else:
        parent[a]=b


n, m = map(int, input().split())
parent = [i for i in range(n+1)]

for _ in range(m):
    k, a, b = map(int, input().split())
    if k==0:
        union(a,b)
    elif k==1:
        # 부모가 같으면 YES 출력하기
        if find(a)==find(b):
            print('YES')
        else:
            print('NO')