[C++] BOJ 1717번: 집합의 표현

학교에서 진행하는 수업에서 공유받은 알고리즘 문제를 풀어보았다.

 

문제의 난이도는 골드 5로, 유니온 파인드 방식을 활용해 푸는 것을 권장받은 문제이다.

문제 링크

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

문제 분석

 이번 문제는 입력으로 n과 m을 받아, n개의 집합을 생성하고 m번의 연산을 수행한다. 연산은 집합을 합하는 합집합과, 두 원소가 같은 집합에 속해 있는지 확인하는 연산 두 가지가 존재한다. 두 원소의 집합을 합하는 합집합은 0 a b의 형태로 입력이 주어지고, 두 원소가 같은 집합에 속해 있는지 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이 때, 두 원소가 같은 집합에 속했는지 확인하는 연산을 수행한 후 같은 집합일 경우엔 "YES"를, 아니라면 "NO"를 출력해야 한다.

자료 구조 및 알고리즘

유니온 파인드

 유니온 파인드는 그래프 알고리즘이다. 두 노드가 같은 그래프에 속하는지 판별할 때에 사용된다. 노드를 합치는 merge(또는 union) 연산과, 노드의 루트 노드를 찾는 find 연산으로 나뉘어진다.

 글로는 설명이 어려울 것 같아 유명한 유튜버분의 소개 영상을 링크로 첨부한다.

  https://youtu.be/AMByrd53PHM?si=S1j9HMk75zJ7lFnT

문제 풀이

 이번 문제를 풀 때에 C++의 stl에 Union Find와 관련하여 구현된 것이 없어 Union Find 방식을 직접 구현해야 하였다. 먼저 처음에는 다음과 같이 parent 배열과 함께 함수를 구현하였다.

 

int parent[1000001]; // 배열의 크기가 10000001인 이유는 N이 1000000을 최대값으로 갖기 때문

int find(int x) {
    if(parent[x] == x) return x;
    return find(parent[x]);
}

void merge(int a, int b) {
    int parentA = find(a);
    int parentB = find(b);
    if (parentA > parentB) 
        parent[parentA] = parentB;
    else
        parent[parentB] = parentA;
}

bool isUnion(int a, int b) {
    return find(a) == find(b);
}

 find 함수는 요소 x가 있을 때, 더 이상 부모가 없는(가장 최상단의) 부모 요소를 return하는 함수이다. 재귀를 통하여 구현하였다.

 merge 함수는 a가 소속된 집합과 b가 소속된 집합을 합치는 함수로, a의 가장 최상단의 부모와 b의 가장 최상단의 부모를 이용하여 합치는 방식으로 구현하였다.

 isUnion 함수는 두 요소가 같은 집합에 속해 있는지 확인하는 함수로, 가장 최상단의 부모가 같다면 같은 집합이므로 find 함수를 활용하여 같은 집합인지 확인하도록 구현하였다.

 

 아래와 같이 main 함수를 작성하고 코드를 제출하였다. 모든 노드의 초기 상태를 표현하기 위해 parent 배열에 대하여 parent[i] = i로 초기화해주었다. 그리고 m번의 연산을 입력받아 수행하였다.

int main(void)
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;

    cin >> n >> m;

    for(int i=1; i<=n; i++)
        parent[i] = i;

    while(m--)
    {
        int op, a, b;
        cin >> op >> a >> b;
        if (op == 0)
            merge(a, b);
        else
            cout << (isUnion(a, b) ? "YES" : "NO") << '\n';
    }

    return 0;
}

 위 코드에서 문제가 발생하였다. 해당 코드는 논리적인 문제는 없었으나, 연산 시간이 오래 걸려 '시간 초과'가 발생하였다. 시간 초과를 해결하기 위해서는 경로 압축이 필요하였다.

 

  경로 압축을 수행하기 위해 find 함수를 아래와 같이 수정하였다.

int find(int x) {
    if(parent[x] == x) return x;
    return (parent[x] = find(parent[x]));
}

이제 모든 자식 노드들은 한 번 find를 수행하고 나면 최상단 부모 노드를 가리키게 되어, find 함수로 다시 탐색하였을 때 소모 시간이 줄어들게 된다.

전체 코드

#include <iostream>

using namespace std;

int parent[1000001]; // 배열의 크기가 10000001인 이유는 N이 1000000을 최대값으로 갖기 때문

int find(int x) {
    if(parent[x] == x) return x;
    return (parent[x] = find(parent[x]));
}

void merge(int a, int b) {
    int parentA = find(a);
    int parentB = find(b);
    if (parentA > parentB) 
        parent[parentA] = parentB;
    else
        parent[parentB] = parentA;
}

bool isUnion(int a, int b) {
    return find(a) == find(b);
}

int main(void)
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;

    cin >> n >> m;

    for(int i=1; i<=n; i++)
        parent[i] = i;

    while(m--)
    {
        int op, a, b;
        cin >> op >> a >> b;
        if (op == 0)
            merge(a, b);
        else
            cout << (isUnion(a, b) ? "YES" : "NO") << '\n';
    }

    return 0;
}