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

📖 백준 1717번 : https://www.acmicpc.net/problem/1717

조건

시간 제한	메모리 제한
2 초	128 MB

문제

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

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

입력

첫째 줄에 $n$, $m$이 주어진다. $m$은 입력으로 주어지는 연산의 개수이다. 다음 $m$개의 줄에는 각각의 연산이 주어진다. 합집합은 $0$ $a$ $b$의 형태로 입력이 주어진다. 이는 $a$가 포함되어 있는 집합과, $b$가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 $1$ $a$ $b$의 형태로 입력이 주어진다. 이는 $a$와 $b$가 같은 집합에 포함되어 있는지를 확인하는 연산이다.

출력

1로 시작하는 입력에 대해서 $a$와 $b$가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

---

풀이

 유니온-파인드 알고리즘을 사용해서 각 노드가 같은 집합에 포함되어 있는지 확인할 수 있다. 문제에서 주의해야 하는 점은 N+1개의 집합이 존재하는 것을 간과해서는 안된다. 따라서 parent배열의 최대 크기를 1,000,001으로 설정해야한다. 또한, N+1개 까지의 원소를 사용할 것이기 때문에 parent 배열을 초기화 할때도 1을 더한 범위까지 초기화해야한다.

코드

#include <iostream>
#include <algorithm>

using namespace std;

int parent[1000001];
int Find(int x) {
    if (parent[x] < 0) return x;//음수 값을 가진다면 x가 루트 노드
    return parent[x] = Find(parent[x]);//경로 압축
}

void Union(int x, int y) {
    x = Find(x);
    y = Find(y);

    if (x == y) return;//이미 같은 집합이므로 무시

    if (parent[x] < parent[y]) {
        parent[x] += parent[y];//x의 크기 증가
        parent[y] = x;//y가 x를 루트로 가리키도록 한다
    }
    else {
        parent[y] += parent[x];//y의 크기 증가
        parent[x] = y;//x가 y를 루트로 가리키도록 한다
    }
}

bool isUnion(int x, int y) {
    x = Find(x);
    y = Find(y);
    return x == y;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);

    int chk, a, b;
    int n, m;
    cin >> n >> m;
    fill(parent, parent + n + 1, -1);//N+1개를 초기화해서 사용할 것

    for (int i = 0; i < m; i++) {
        cin >> chk >> a >> b;
        if (!chk) Union(a, b);
        else isUnion(a, b) ? cout << "YES" << '\n' : cout << "NO" << '\n';
    }
    return 0;
}