[백준] 집합의 표현(1717) - python

✏️ 문제

초기에 {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 를 출력해도 된다)

🎈 입출력 예

<입력>
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

<출력>
NO
NO
YES

👩‍💻 내 코드

이 문제는 Union Find의 완전 기본 문제이다. 따로 아이디어 구상 없이 바로 알고리즘을 적용해서 풀었다.

import sys
sys.setrecursionlimit(10**6)


def find(node):
    if node == parent[node]:
        return node
    parent[node] = find(parent[node])
    return parent[node]


def union(a, b):
    pa = find(a)
    pb = find(b)

    if pa == pb:
        return

    parent[pb] = pa


if __name__ == "__main__":

    n, m = map(int, input().split())
    parent = [i for i in range(n+1)]

    for _ in range(m):
        cmd, a, b = map(int, sys.stdin.readline().split())

        if cmd == 0:
            union(a, b)
        else:
            if find(a) == find(b):
                print("YES")
            else:
                print("NO")