[백준] 집합의 표현 1717번 - Python

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

1717번: 집합의 표현

문제

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

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

입력

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

출력

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

---

풀이

유니온 파인드 ( union - find )는 그래프 알고리즘이다.

두 집합을 하나의 집합으로 묶는 union 연산과 원소 x가 속한 집합의 대표 노드를 찾는 find 연산으로 구성되어있는 알고리즘이다.

 

문제의 입력에서 0 a b 형태는 합집합, union 연산을 해주어야하고 1 a b 형태는 원소 a와 b에 대해 find 연산을 하여 대표노드들을 찾은 후 같은지(두 원소가 같은 집합에 포함되어있는지) 확인해주는 연산이다.

import sys
input=sys.stdin.readline
sys.setrecursionlimit(1000000) # 재귀 깊이 제한 늘리기

n,m = map(int, input().split())

# n개 노드들의 대표노드 배열
s = [0]*(n+1)
for i in range(n+1):
    s[i] = i

#대표노드 반환, 경로압축 포함
def find(i): 
    if s[i] != i:
        s[i] = find(s[i])
    return s[i]


def union(a,b): 
    a=find(a) #a가 속한 집합의 대표노드
    b=find(b) #b가 속한 집합의 대표노드

    if a > b: # 더 작은 것을 대표노드로 하여 합집합 수행
        s[a] = s[b]
    elif a < b:
        s[b] = s[a]

for _ in range(m):
    q, a, b=list(map(int, input().split()))
    if q == 0:
        union(a, b)
    elif q == 1:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")

 

find에서의 경로단축 관련.