algorithm/small_to_large.cpp at main · CAEL0/algorithm · GitHub

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
91
92
93
94
95
96
97
98
99
100
101
102
// BOJ 17469 트리의 색깔과 쿼리

#include <bits/stdc++.h>
#define sz size()
#define bk back()
#define fi first
#define se second

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;

struct DisjointSet {
    int n;
    vector<int> dsu, rank;

    DisjointSet(int _n) {
        n = _n;
        dsu.resize(n + 1);
        rank.resize(n + 1);
        init();
    }

    void init() {
        for (int i = 1; i <= n; i++) {
            dsu[i] = i;
            rank[i] = 1;
        }
    }

    int find(int z) {
        if (z != dsu[z])
            dsu[z] = find(dsu[z]);
        return dsu[z];
    }

    void merge(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y)
            return;

        if (rank[x] < rank[y])
            swap(x, y);

        rank[x] += rank[y];
        dsu[y] = x;
    }

    bool is_same(int x, int y) { return find(x) == find(y); }
};

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

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

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

    vector<unordered_set<int>> colors(n + 1);
    for (int i = 1; i <= n; i++) {
        int x;
        cin >> x;

        colors[i].insert(x);
    }

    q += n - 1;
    vector<pii> query(q);
    for (int i = 0; i < q; i++)
        cin >> query[i].fi >> query[i].se;

    DisjointSet ds(n);
    ds.init();

    vector<int> ans;
    for (int i = q - 1; i >= 0; i--) {
        int op = query[i].fi;
        int x = query[i].se;

        if (op == 1) {
            int y = ds.find(parent[x]);
            x = ds.find(x);

            if (colors[x].sz < colors[y].sz)
                swap(x, y);

            ds.dsu[y] = x;
            for (int z : colors[y])
                colors[x].insert(z);
        } else if (op == 2)
            ans.push_back(colors[ds.find(x)].sz);
    }

    for (int i = ans.sz - 1; i >= 0; i--)
        cout << ans[i] << '\n';
}