BOJ :: 집합의 표현 (no.1717)

문제

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

🤔 생각

• 그래프로 치면 부모 노드를 계속 찾아야 한다는 소리다.
• 유니온 파인드로 풀 수 있겠다.

📌 내 풀이

# 집합의 표현
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**9)

def find(x):
    if p[x] == x: return x
    p[x] = find(p[x])
    return p[x]

def union(a,b):
    x = find(a)
    y = find(b)

    if x == y: return
    if x < y: p[x] = y
    else: p[y] = x

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

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

• 한가지 주의할 점은, find 함수에 재귀 결과를 부모 배열에 저장하지 않으면 시간 복잡도가 훨씬 커진다.
• 물론 안 해도 답은 나오지만, 어차피 가장 루트 노드만 비교하면 되므로 압축한 것이라 생각하면 되겠다.

✔ 회고

• 유니온 파인드 재활 완료!!
• 더 익숙해지기 위해 쭉쭉 밀자