[Python] 백준 1717번 - 집합의 표현 (골드 5)

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

1717번: 집합의 표현

 

혼자 힘으로 풀었는가? X

알고리즘 분류
 - 자료 구조
 - 분리 집합

 

문제

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

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

입력

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

출력

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

제한

• $1 ≤ n ≤ 1\,000\,000$ 
•  $1 ≤ m ≤ 100\,000$ 
•  $0 ≤ a, b ≤ n$ 
•  $a$, $b$는 정수
•  $a$와 $b$는 같을 수도 있다.

 

---

 

2023.09.05 - [Language/Python] - 분리 집합 (Disjoint Set) : Union-Find 알고리즘

분리 집합 (Disjoint Set) : Union-Find 알고리즘

 

위의 알고리즘을 사용한다.

 

import sys
input = sys.stdin.readline
sys.setrecursionlimit(100000)
n, m = map(int, input().split())

parent = [i for i in range(n+1)]

#주어진 집합의 루트를 찾음
def find(x):
    if parent[x] == x: #주어진 값이 루트면 리턴
        return x
    parent[x] = find(parent[x]) #아니면 부모를 찾아 재귀
    return parent[x]

#두 개의 집합이 같은 집합인지 확인, 아닐경우 하나로 합침
def union(x, y):
    x = find(x)
    y = find(y) #주어진 x, y의 루트를 찾아서
    if x == y: #루트가 같다면 동일한 집합
        return
    parent[x] = y #아니라면 같은 집합으로 합쳐주기
    #x의 부모는 y
    return

for _ in range(m):
    flag, a, b = map(int, input().split())
    if flag == 0:
        union(a, b)
    else:
        x = find(a)
        y = find(b)
        if x == y:
            print("YES")
        else:
            print("NO")

parent 배열의 루트를 초기값을 자기자신으로 할당한다.

 

parent = [i for i in range(n+1)]

이후 1이 입력될 때 a와 b의 루트를 찾아서 동일한 루트(같은 집합)이면 YES, 아니면 NO를 출력한다.

x = find(a)
y = find(b)
if x == y:
    print("YES")
else:
    print("NO")