백준 온라인 저지, 분리집합 / 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

힌트

접근 방법

- 유니온 파인드를 사용해 같은 집합에 속하는지를 확인한다.
- 유니온 파인드
- 0. n+1개만큼 길이륵 가지는 리스트를 만들어 각각의 인덱스에 해당하는 값을 넣어 초기화한다.
- 1. 재귀함수를 활용해 부모노드를 구하는 함수를 만든다. 해당 함수는 인덱스와 값이 같은 경우에 최종적인 부모노드를 리턴하는 함수이다.
- 2. 합집합을 만드는 함수를 만든다. 해당 함수는 두 집합의 최종적인 부모노드의 값을 비교해 더 작은 값이 부모가 되도록(더 큰 값의 인덱스에 작은 값을 저장) 구현한다.
- 3. 같은 집합에 속해있는지 판단하는 함수를 만든다. 해당함수는 부모노드를 찾아 두 원소의 부모노드가 같다면 True, 아니면 False를 출력하도록한다.

코드

# 접근 방법
# 유니온 파인드를 사용해 같은 집합에 속하는지를 확인한다.
# 유니온 파인드
# 0. n+1개만큼 길이륵 가지는 리스트를 만들어 각각의 인덱스에 해당하는 값을 넣어 초기화한다.
# 1. 재귀함수를 활용해 부모노드를 구하는 함수를 만든다. 해당 함수는 인덱스와 값이 같은 경우에 최종적인 부모노드를 리턴하는 함수이다.
# 2. 합집합을 만드는 함수를 만든다. 해당 함수는 두 집합의 최종적인 부모노드의 값을 비교해 더 작은 값이 부모가 되도록(더 큰 값의 인덱스에 작은 값을 저장) 구현한다.
# 3. 같은 집합에 속해있는지 판단하는 함수를 만든다. 해당함수는 부모노드를 찾아 두 원소의 부모노드가 같다면 True, 아니면 False를 출력하도록한다.
def get_parent(index):
    if parent[index] == index:
        return index
    parent[index] = get_parent(parent[index])
    return parent[index]

def union_parent(a, b):
    x1 = get_parent(a)
    x2 = get_parent(b)
    parent[max(x1, x2)] = min(x1, x2)

def find_parent(a, b):
    x1 = get_parent(a)
    x2 = get_parent(b)
    if x1 != x2:
        return print('NO')
    else:
        return print('YES')

import sys
sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline
n, m = map(int, input().split())
parent = [i for i in range(n+1)]
for _ in range(m):
    check, a, b = map(int, input().split())
    if check:
        find_parent(a, b)
    else:
        union_parent(a, b)