[백준] 1717번. 집합의 표현 (Python 파이썬)

문제

시간 제한메모리 제한제출정답맞힌 사람정답 비율

2 초

128 MB

92682

29405

17922

28.124%

문제

초기에 �+1$n+1$개의 집합 {0},{1},{2},…,{�}$\{0\}, \{1\}, \{2\}, \dots , \{n\}$이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

입력

첫째 줄에 �$n$, �$m$이 주어진다. �$m$은 입력으로 주어지는 연산의 개수이다. 다음 �$m$개의 줄에는 각각의 연산이 주어진다. 합집합은 0$0$ �$a$ �$b$의 형태로 입력이 주어진다. 이는 �$a$가 포함되어 있는 집합과, �$b$가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1$1$ �$a$ �$b$의 형태로 입력이 주어진다. 이는 �$a$와 �$b$가 같은 집합에 포함되어 있는지를 확인하는 연산이다.

출력

1로 시작하는 입력에 대해서 �$a$와 �$b$가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

제한

•  1≤�≤1000000$1 ≤ n ≤ 1\,000\,000$ 
•  1≤�≤100000$1 ≤ m ≤ 100\,000$ 
•  0≤�,�≤�$0 ≤ a, b ≤ n$ 
•  �$a$, �$b$는 정수
•  �$a$와 �$b$는 같을 수도 있다.

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

풀이

어렵진 않다. 서로소 집합 (Disjoint set) 의 기본적인 이론만 알면 간단히 풀 수 있는 문제다. 

그치만, 시간초과와 런타임 에러 (RecursionError) 가 충분히 날 수 있는 문제다. 

 

나는 sys.stdin.readline 으로 10만 (m) 이나 되는 입력값을 빠르게 받았고, 

setrecursionlimit 를 10 의 8승으로 올려 런타임에러를 잡았다. 

 

import sys

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

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

input = sys.stdin.readline

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

for _ in range(m):
    oper,a,b = map(int, input().split())

    if oper == 0:
        union_parent(parent,a,b)
    else:
        if find_parent(parent,a) == find_parent(parent,b):
            print('YES')
        else:
            print('NO')