[Baekjoon/C++] 1717번 - 집합의 표현

[백준] Baekjoon Online Judge

문제로 이동

 

문제

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

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

 

입력

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

 

출력

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

 

---

풀이

#include <iostream>
using namespace std;

int n, m;
int parent[1000001];

int findParent(int n);
bool isSame(int n, int a, int b);

int findParent(int n) {
    if (parent[n] == n) return n;

    return parent[n] = findParent(parent[n]);
}

bool isSame(int n, int a, int b) {
    int aa = findParent(a);
    int bb = findParent(b);

    if (aa == bb) return true;
    else {
        if (n == 0) parent[bb] = aa;

        return false;
    }
}

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

    cin >> n >> m;

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

    while (m--) {
        int n, a, b;
        cin >> n >> a >> b;

        bool same = isSame(n, a, b);

        if (n == 1) {
            if (same) cout << "YES\n";
            else cout << "NO\n";
        }
    }

    return 0;
}

 

int findParent(int n) {
    if (parent[n] == n) return n;

    return parent[n] = findParent(parent[n]);
}

 배열을 탐색하며 n의 루트 노드를 찾는 함수이다.

 

bool isSame(int n, int a, int b) {
    int aa = findParent(a);
    int bb = findParent(b);

    if (aa == bb) return true;
    else {
        if (n == 0) parent[bb] = aa;

        return false;
    }
}

 입력 받은 두 원소가 같은 집합에 포함되어 있는지 확인하는 함수이다. findParent()를 통해 각 원소의 루트 노드를 찾고, 그 값이 서로 같다면 두 원소는 같은 집합의 원소인 것이다. 만약 두 원소가 같은 집합이 아니고 n이 0이라면 두 집합을 합집합 연산을 하도록 했다.