[백준/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 ≤ n ≤ 1\,000\,000$  
• $1 ≤ m ≤ 100\,000$  
• $0 ≤ a, b ≤ n$ 
• $a$, $b$는 정수
• $a$와 $b$는 같을 수도 있다.

풀이
오늘 공부한 유니온 파인드 알고리즘을 활용해볼 수 있는 정석적인 문제였다
union 함수와 find 함수를 각각 작성한 후에 for 루프를 m번 돌면서 명령에 맞게 union과 find를 호출해준다
명령이 1인 경우에는 a에 대한 find 함수의 결과와 b에 대한 find 함수의 결과를 비교해 답을 출력한다
주의할 점은 재귀함수를 쓸 때는 무조건 sys.setrecursionlimit()을 설정하자

코드

import sys
sys.setrecursionlimit(1000000)

n, m = map(int, sys.stdin.readline().split())

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

def find_parent(parent, a):
    if parent[a] != a:
        parent[a] = find_parent(parent, parent[a])
    return parent[a]

def union_node(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


for _ in range(m):
    command, a, b = map(int, sys.stdin.readline().split())
    if command == 0:
        union_node(parent, a, b)
    else:
        print("yes" if find_parent(parent, a) == find_parent(parent, b) else "no")