백준 1717 집합의 표현 (C++)

1717번: 집합의 표현

dsu를 활용한 문제이다. find-union 방식을 알고 있다면 어렵지 않게 풀 수 있는 문제이다. 문제에서 주어지는 집합을 합치는 경우는 루트를 union하는 값으로 바꿔주는 방식으로 저장하였다. 만약 두 값이 같은 집합 내에 있다면, 두 값은 같은 루트를 가지게 될 것 이므로 루트가 같을 경우 YES, 아니면 NO를 출력해주었다.
바로 바로 아이디어가 생각나지 않았다. dsu 문제를 많이 풀어봐야겠다.

#include <iostream>
#include <vector>

using namespace std;

int n, m;
vector<pair<int, pair<int, int>>> v;
int root[1000001];

int findRoot(int a) {
    if (root[a] == a) return a;
    else return root[a] = findRoot(root[a]);
}

void unionRoot(int a, int b) {
    a = findRoot(a);
    b = findRoot(b);

    if (a != b) root[a] = b;
}

void solution() {
    for (int i = 0; i <= n; i++) {
        root[i] = i;
    }

    for (int i = 0; i < m; i++) {
        int check = v[i].first;
        int a = v[i].second.first;
        int b = v[i].second.second;

        if (check == 1) {
            if (findRoot(a) == findRoot(b)) cout << "YES\n";
            else cout << "NO\n";
        }
        else unionRoot(a, b);
    }
}

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

    cin >> n >> m;

    for (int i = 0; i < m; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        v.push_back({ a,{b,c} });
    }

    solution();

    return 0;
}