[백준 Python] 1717번 집합의 표현

문제

초기에 $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. 연산의 맨 앞이 0인 경우에는 두 원소를 같은 집합에 넣기
# 2. 연산의 맨 앞이 1인 경우에는 같은 집합에 있는지를 확인하기

def find_parent(x):
    # 초기 상태에서 find_parent(x)하면 x값이 그대로 나와야 함
    # 루트 노드를 찾아갈 수 있도록 재귀적으로 작성
    if data[x] != x:
        data[x] = find_parent(data[x])
    return data[x]

def union_subset(x,y):
    # x와 y원소를 받아서 그 두 개를 합치기
    # x가 속한 집합과 y가 속한 집합을 찾아서 그 둘을 합치는 것
    x = find_parent(x)
    y = find_parent(y)
    # 만약 둘이 속한 집합이 다르면 합쳐주기 -> 같은 부모노드를 가지고 있다고 처리
    if x != y:
        data[x] = y

# 시간초과 방지를 위한 처리
import sys
input = sys.stdin.readline
sys.setrecursionlimit(100000)

# n+1개의 집합이 존재하므로 range(n+1)
# 현재 인덱스에 저장된 값은 자신의 부모노드인 셈
n,m = map(int,input().split())
data = [i for i in range(n+1)]

for _ in range(m):
    part, x, y = map(int,input().split())
    if part == 0:
        union_subset(x,y)
    if part == 1:
        find_x = find_parent(x)
        find_y = find_parent(y)
        if find_x != find_y :
            print("NO")
        if find_x == find_y:
            print("YES")