[BOJ] 13164 행복 유치원

문제

1717번: 집합의 표현

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

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

풀이

• find(a) : a가 속한 집합의 루트 노드를 찾는 연산 경로 압축을 사용하여 find를 하면서 만난 모든 값의 부모 노드를 root로 만든다.

• merge(a, b) : a가 속한 집합과 b가 속한 집합을 합치는 연산 h에 트리의 높이를 저장하여 항상 높이가 더 낮은 트리를 높은 트리 밑에 넣는다. 만약 두 트리의 높이가 같다면 +1을 해준다. cf) (루트노드(1) + 원래 트리의 높이(h[a])) 이므로

코드

#include <cstdio>
#include <cstring>
using namespace std;
 
int n, m;
int p[1000001]; // 루트 노드 저장
int h[1000001]; // 트리의 높이
 
int find(int a) {
    if (p[a] < 0) return a;
    return p[a] = find(p[a]);
}
 
void merge(int a, int b) {
    a = find(a), b = find(b);
    if (a == b) return; // 두 값의 root가 같으면 합치지 않는다.
 
    // 항상 높이가 낮은 트리를 높이가 큰 트리 밑에 넣는다.
    // 즉, 높이가 더 큰 트리가 루트가 된다.
    if (h[a] < h[b]) p[a] = b;
    else {
        p[b] = a;
        if (h[a] == h[b]) h[a]++;
    }
}
 
int main() {
    memset(p, -1, sizeof(p));
    memset(h, 1, sizeof(h));
    scanf("%d %d", &n, &m);
    int a, b, c;
    while (m--) {
        scanf("%d %d %d", &a, &b, &c);
        if (!a) merge(b, c);
        else {
            if (find(b) == find(c)) printf("YES\n");
            else printf("NO\n");
        }
    }
}