[백준 1717번 ] 파이썬 - 집합의 표현

백준 1717 - 집합의 표현

시간 제한 2초, 메모리 제한 128MB

# 조건

• 초기에 $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를 이용하여 풀어준다.
• 속한 집합을 나타낼 parents 리스트를 [i for i in range(N)]으로 만들어준다.
• 0번 명령어가 들어왔다면
  • a와 b를 union 시켜주는데 작은 번호를 기준으로 해준다.
• 1번 명령어가 들어왔다면 각각의 find 값이 같다면 Yes, 아니라면 No를 출력해준다.
  • 여기서 parent 값을 바로 확인하지 않는 이유는
  • 1 2, 3 4 의 집합을 합친 후 2 3 을 합친다면
  • 4번의 집합 3번으로 기록되어 있기 때문이다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  

def find(x):  
    if parents[x] != x:  
       x = find(parents[x])  
    return x  

def union(x, y):  
    x = find(x)  
    y = find(y)  
    if x > y:  
        parents[x] = y  
    else:  
        parents[y] = x  


N, M = map(int, input().split())  
parents = [i for i in range(N+1)]  
for _ in range(M):  
    q, a, b = map(int, input().split())  

    if q == 0:  
        union(a, b)  
    else:  
        print('YES' if parents[a] == parents[b] or find(a) == find(b) else 'NO')