洛谷U639895 树上子树颜色询问2 题解 树上启发式合并（dsu on tree）

原题出处

题目链接：https://www.luogu.com.cn/problem/U639895

解题思路：

基本上和第一题一样。树上启发式合并。

示例程序：

#include <bits/stdc++.h>
using namespace std;
const int maxn = 2e5 + 5;

int n, m, c[maxn], t[maxn], ans[maxn], sz[maxn], son[maxn];
vector<int> g[maxn];

struct Query {
    int id, c;
};
vector<Query> query[maxn];

void getsz(int u, int p) {
    sz[u] = 1;
    for (auto v : g[u]) {
        if (v != p) {
            getsz(v, u);
            sz[u] += sz[v];
            if (sz[v] > sz[ son[u] ])
                son[u] = v;
        }
    }
}

void dfs_clean(int u, int p) {
    t[ c[u] ]--;
    for (auto v : g[u])
        if (v != p)
            dfs_clean(v, u);
}

void dfs(int u, int p, bool keep) {
    for (auto v : g[u])
        if (v != p && v != son[u])
            dfs(v, u, false);
    if (son[u])
        dfs(son[u], u, true);
    for (auto v : g[u])
        if (v != p && v != son[u])
            dfs(v, u, true);
    t[ c[u] ]++;
    for (auto &[id, c] : query[u])
        ans[id] = t[c];
    query[u].clear();
    if (!keep)
        dfs_clean(u, p);
}

int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++)
        scanf("%d", c+i);
    for (int i = 1, u, v; i < n; i++) {
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    scanf("%d", &m);
    for (int i = 1, u, c; i <= m; i++) {
        scanf("%d%d", &u, &c);
        query[u].push_back({i, c});
    }
    getsz(1, -1);
    dfs(1, -1, true);
    for (int i = 1; i <= m; i++)
        printf("%d\n", ans[i]);
    return 0;
}