백준 17469번: 트리의 색깔과 쿼리 [C++]

작성일 | In
플래티넘 III 플래티넘 III

문제

트리의 색깔과 쿼리

풀이

쿼리를 마지막부터 거꾸로 처리하여 간선을 제거하는 쿼리가 아니라 간선을 추가하는 쿼리로 생각해 봅시다.

분리 집합을 이용하여 분리 집합 트리의 루트마다 색깔들을 담은 집합을 하나씩 저장해 줍시다. 간선을 추가할 때 노드 뿐 아니라 색라 집합도 함께 합쳐 줍시다.

작은 집합에서 큰 집합으로 합치는 테크닉을 이용하면 집합의 각 원소가 $\mathcal{O}(\log N)$번 이동하므로 해시를 이용한 집합을 사용하면 시간 복잡도 $\mathcal{O}(N\log N)$에 해결할 수 있습니다.

코드

코드에서는 트리를 이용한 집합을 사용하였으므로 시간 복잡도는 $\mathcal{O}(N\log^2N)$입니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#include <bits/stdc++.h>
using namespace std;

struct DisjointSet {
    vector<set<int>> sets;
    vector<int> parent;

    DisjointSet(size_t size): sets(size), parent(size) {
        for (int i = 0; i < (int)size; i++)
            parent[i] = i;
    }

    int findRoot(int x) {
        if (x != parent[x])
            parent[x] = findRoot(parent[x]);
        return parent[x];
    }

    void unionRoot(int u, int v) {
        u = findRoot(u);
        v = findRoot(v);

        if (sets[u].size() < sets[v].size())
            swap(u, v);

        for (int e : sets[v])
            sets[u].insert(e);

        sets[v].clear();
        parent[v] = u;
    }
};

struct Query {
    int type, a;
    int answer;
} queries[1111111];

int parent[101010];

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

    int N, Q;
    cin >> N >> Q;

    DisjointSet disjointSet(N + 1);

    int totalQuery = N + Q - 1;

    for (int i = 2; i <= N; i++)
        cin >> parent[i];
    for (int i = 1; i <= N; i++) {
        int c;
        cin >> c;
        disjointSet.sets[i].insert(c);
    }

    vector<bool> removed(N + 1);
    for (int i = 1; i <= totalQuery; i++) {
        auto& [type, a, _] = queries[i];
        cin >> type >> a;
        if (type == 1)
            removed[a] = true;
    }

    for (int i = 2; i <= N; i++)
        if (not removed[i])
            disjointSet.unionRoot(i, parent[i]);
    
    for (int i = totalQuery; i > 0; i--) {
        auto& [type, a, answer] = queries[i];

        if (type == 1)
            disjointSet.unionRoot(a, parent[a]);
        else {
            int r = disjointSet.findRoot(a);
            answer = disjointSet.sets[r].size();
        }
    }

    for (int i = 1; i <= totalQuery; i++) {
        auto& [type, _, answer] = queries[i];
        if (type == 2)
            cout << answer << '\n';
    }

    return 0;
}

백준 17469번: 트리의 색깔과 쿼리 [C++]