백준 #1717 집합의 표현 (파이썬) : Union-Find(합집합 찾기)

오늘의 한 마디
Class 5 깨기 시작!

• Baekjoon 문제 링크
• solved.ac > 문제 > CLASS 5에 수록된 문제입니다.

문제

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

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

그 유명한 합집합 찾기(Union-Find) 알고리즘을 사용하는 문제들을 분리 집합이라는 알고리즘 분류로 묶은 것이었다.

나동빈 님의 Union-Find 설명이 가장 깔끔하다.

• 유튜브 영상
• 블로그 글

약간 헷갈릴 수 있는 것이, 사실 집합 내 원소 간의 부모-자식 관계는 없다.

같은 집합 내의 원소인가? 라는 문제를 공통의 부모 원소를 갖는가? 라는 문제로 환원하기 위해서 가상의 부모를 상정했을 뿐이다.

위 글의 Union-Find 알고리즘에서는 합칠 때, 더 작은 부모의 원소 쪽으로 합친다.

풀이

이 풀이는 나동빈 님의 C++ 풀이를 참고하여 Python으로 구현한 것입니다.

import sys
sys.setrecursionlimit(10**5)

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

parent = [i for i in range(n+1)]

def get_parent(x):
    if parent[x] == x:
        return x
    parent[x] = get_parent(parent[x]) # get_parent 거슬러 올라가면서 parent[x] 값도 갱신
    return parent[x]

def union_parent(a, b):
    a = get_parent(a)
    b = get_parent(b)

    if a < b: # 작은 쪽이 부모가 된다. (한 집합 관계라서 부모가 따로 있는 건 아님)
        parent[b] = a
    else:
        parent[a] = b        

def find_parent(a, b):
    return get_parent(a) == get_parent(b)

for _ in range(m):
    op, a, b = map(int, sys.stdin.readline().rstrip().split())
    if op == 0:
        union_parent(a, b)
    else:
        print('YES' if find_parent(a, b) else 'NO')