백준 1717 집합의 표현 (파이썬)

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

1717번: 집합의 표현

1. 집합 n+1개를 만들고 자기자신을 부모로 설정한다

2. z가 0이면 앞의 숫자를 부모로 하는 disjoint-set 실행

3. z가 1이면 각자의 부모를 찾아 일치하면 YES 아니면 NO 출력

import sys

input = sys.stdin.readline

def get_parent(x):
    if parent[x] == x:
        return x
    p = get_parent(parent[x])
    parent[x] = p
    return p

def union(x, y):
    x = get_parent(x)
    y = get_parent(y)

    if x != y:
        parent[y] = x

def find_parent(x):
    if parent[x] == x:
        return x
    return find_parent(parent[x])

n, m = map(int, input().split())
parent = {}

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

for _ in range(m):
    z, x, y = map(int, input().split())

    if not z:
        union(x, y)

    if z:
        if find_parent(x) == find_parent(y):
            print('YES')
        else:
            print('NO')