[BOJ 13306 | C++] 트리 :: Union-Find, 오프라인 쿼리

• 문제 링크 : https://www.acmicpc.net/problem/13306
• 문제 태그 : 자료 구조, 분리 집합, 오프라인 쿼리
• 관련 문제 : https://justicehui.github.io/medium-algorithm/2024/02/04/union-find-application/

Union Find 200% 활용하기

 

풀이

• 비재귀를 이용한 find함수를 구현하고 경로 압축을 적용했습니다.
• 스택을 이용하여 쿼리를 저장하고 역순으로 쿼리를 수행했습니다.
• 정점 c와 d를 연결하는 경로가 존재하는지 확인하는 쿼리는 두 정점의 부모가 같은지 비교하는 방식으로 처리했습니다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <stack>
#include <tuple>

using namespace std;

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

    int n, q;
    cin >> n >> q;

    vector<int> p(n + 1, -1);

    auto find = [&](int x) {
        while (p[x] >= 0) x = p[x];
        return x;
    };

    auto merge = [&](int u, int v) -> bool{
        u = find(u);
        v = find(v);
        if (u == v) return 0;
        if (-p[u] < -p[v]) swap(u, v);

        p[u] += p[v];
        p[v] = u;
        return 1;
    };

    vector<int> parent(n + 1);
    for (int i = 1; i < n; i++) {
        int x;
        cin >> x;
        parent[i + 1] = x;
    }

    stack<tuple<int, int, int>> query;
    stack<string> ans;

    for (int i = 0; i < n - 1 + q; i++) {
        int x, u, v;
        cin >> x;
        if (x == 0) {
            cin >> u;
            query.push({ 0, u, parent[u] });
        }
        else {
            cin >> u >> v;
            query.push({ 1, u, v });
        }
    }

    while (!query.empty()) {
        auto [x, u, v] = query.top();
        query.pop();
        if (x == 0) {
            merge(u, v);
        }
        else {
            if (find(u) == find(v)) ans.push("YES\n");
            else ans.push("NO\n");
        }
    }

    while (!ans.empty()) {
        cout << ans.top();
        ans.pop();
    }
}