[Python - Data Structure, Union Find] 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 를 출력해도 된다)

📌 문제 풀이

👨‍🏫 접근

합치기 연산 알고리즘으로 집합을 판별하는 전형적인 알고리즘이다.

 

문제처럼 집합이 주어졌을 때, 이 집합이 서로소인지 아닌지 판별하고 그 집합을 합치기 위해 사용한다. 먼저 "부모"라는 것을 설정해주도록 한다.

 

예를 들어, {1}, {2}, {3}, {4}, {5}, {6}, {7} 이라는 집합들이 있다고 할 때, 각각의 부모를 자기 자신으로 설정하는 것이다. 그래서 [1, 2, 3, 4, 5, 6, 7]  이라는 부모들이 존재한다.

 

이제 합치기 연산에서 이 부모가 큰 역할을 한다. {1}과 {3} 집합을 합한다고 하면, {1, 3}으로 간단하게 연산할 수 있지만, 코드상으로 구현하기 단순하지 않다. 그래서 이 부모 리스트를 통해 집합을 표현하는 것이다. 1과 3을 합친다고 한다면, 둘 중에 더 작은 값을 부모로 설정하는 것이다. 어떤 값을 부모로 설정할지는 중요하지 않아서 더 큰 값을 부모로 설정해도 된다.

 

각설하고 위의 연산을 마치면, [1, 2, 1, 4, 5, 6, 7] 이라는 부모 리스트가 만들어진다. 즉 집합의 세 번째 요소는 1을 부모로 한다는 것이다. 이 의미는 1과 3이 같은 집합에 속함임을 말해준다.

 

동일하게 4와 7을 합해준다. 그러면 [1, 2, 1, 4, 5, 6, 4]가 된다.

 

이제 3과 7을 합치는 연산을 해보자. 그러면 최종적으로 {1, 3, 4, 7}이라는 집합이 만들어진다. 이때 이를 구현하기 위해서는 3을 포함하는 집합과 7을 포함하는 집합의 부모를 찾아서 그 부모들끼리 다시 연산을 하는 것이다.

3의 부모는 1이며, 7의 부모는 4이다. 그래서 1과 4 중에 부모를 정한다고 하면 1이 부모가 되므로, 4의 부모는 1이 된다. 즉 최종적으로 [1, 2, 1, 1, 5, 6, 4] 의 부모 리스트가 만들어진다. 7번째 요소는 어차피 4를 부모로 가지므로, 4의 부모를 찾도록 연산하면 7도 {1, 3, 4, 7}에 속한다는 사실을 알 수 있다.

👨‍🏫 문제 풀이

📄 전체 코드

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline


def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]


def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


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

for i in range(m):
    com, a, b = map(int, input().split())
    if com == 0:
        union_parent(parent, a, b)
    else:
        if find_parent(parent, a) != find_parent(parent, b):
            print("NO")
        else:
            print("YES")

    print(parent)

📄 준비

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline


def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]


def union_parent(parent, a, b):
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


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

부모를 찾는 연산을 해준다. 경로압축을 통해서 빠르게 부모를 찾아낸다.

 

그 이후로는 부모를 합치는 연산을 해준다. 이때, 각 요소들의 최종 조상을 찾아내서 합쳐준다.

 

부모는 자기자신으로 초기화해준다.

📄 풀이

for i in range(m):
    com, a, b = map(int, input().split())
    if com == 0:
        union_parent(parent, a, b)
    else:
        if find_parent(parent, a) != find_parent(parent, b):
            print("NO")
        else:
            print("YES")

    print(parent)

합치는 연산에서는 합쳐주고, 같은 집합인지 찾는 연산에서는 최종 조상을 찾는 연산을 해주어 같을 경우에만 YES를 출력하고 아니면 NO를 출력시킨다.

📌 총평

분리집합 연산을 까먹고 있었는데 이번 기회에 복습해서 기분이 좋다.