[Python] 백준 1717: 집합의 표현 - disjoint set (Union find)

문제

1717번: 집합의 표현

풀이

처음엔 문제에서 설명하는 순서대로 인접리스트를 만들고, set으로 풀었는데 메모리 초과가 났다.

따라서 disjoint set을 사용해서 풀어야하는데

루트배열을 두고 서로 루트가 같으면 같은 집합에 속해있는것으로 판단하면 된다.

먼저 n의 크기의 root 배열을 만들고, i(본인)로 채운다.

i	0	1	2	3	4	5	6	7
root[i]	0	1	2	3	4	5	6	7

 

 

그 후 a, b 두 집합을 합치기 위해서는 a의 루트를 b에 연결시킨다.

예를들어 (부모)1-2-3-4 이렇게 연결된 집합 a(1)와 (부모)5-6-7 이렇게 연결된 집합 b(5)가 있을때, a를 b에 연결시키면

5-1-2-3-4

  -6-7

이런 모양이 된다. 그러면 1~7까지 모두 하나의 집합이 되고 모두 루트는 5가된다. 

i	0	1	2	3	4	5	6	7
root[i]	0	5	1	2	3	5	5	6

 

그런데 루트배열에 저장된 값은 부모이긴 하지만 최상위 루트라는 보장이 없기 때문에 최상위루트를 찾을 수 있는 find함수를 만들어줬다. 

 

 

import sys

n, m = map(int, sys.stdin.readline().split(" "))

# n은 a, b가 될 수 있는 최대값, 0은 비워두기
root = [i for i in range(n+1)]

# 최상위 루트 찾기
def find(a):
    if root[a] != a:
        root[a] = find(root[a])
    return root[a]

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

    # 두 원소가 같은 집합에 포함되어 있는지를 확인 (YES/NO 출력)
    if cmd == 1:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")

    # b에 a의 최상위 루트 붙이기
    else:
        if find(a) != find(b):
            root[find(a)] = b