문제

초기에 $n+1$개의 집합 $\{0\}, \{1\}, \{2\}, \dots , \{n\}$이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

문제 링크

풀이

union-find 알고리즘을 먼저 이해하고 이 문제를 본다면 쉽게 풀 수 있다.
연결 관계를 트리로 구현했고 각 노드의 트리 높이를 depth라는 배열에 저장하여 최적화를 하였다.

a와 b가 같은 집합인지 확인하는 방법은
각 요소의 루트 즉 find(a)와 find(b)를 비교해보면 된다.

코드

#include <iostream>
#include <cstring>

using namespace std;

int n, m;
int *ptr;
int *depth;


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

void merge(int a, int b){
    a = find(a);
    b = find(b);

    if (a == b) return;

    if (depth[a] > depth[b]){
        ptr[b] = a;
    }
    else if(depth[a] < depth[b]){
        ptr[a] = b;
    }
    else{
        ptr[a] = b;
        depth[b]++;
    }


}

int main() {
    scanf("%d %d", &n, &m);
    ptr = (int*)malloc( sizeof(int) * n+1 );
    depth = (int*)malloc( sizeof(int) * n+1 );
    memset(depth, 0, sizeof(int) * n+1 );
    for(int i=0; i<=n; i++){
        ptr[i] = i;
    }

    for (int i =0; i < m; i++){
        int op, a, b;
        // cin >> op >> a >> b;
        scanf("%d %d %d", &op, &a, &b);
        switch(op){
            case 0:
                merge(a, b);
                break;

            case 1:
                if (find(a) == find(b)) printf("YES\n");
                else printf("NO\n");
                break;
        }
    }

    free(ptr);
    return 0;
}

참고

https://gmlwjd9405.github.io/2018/08/31/algorithm-union-find.html