【板子】树上启发式合并

给出一棵 𝑛 个节点以 1 为根的树，节点 𝑢 的颜色为 𝑐𝑢，现在对于每个结点 𝑢 询问以 𝑢 为根的子树里一共出现了多少种不同的颜色．

支持离线，考虑预处理后 𝑂(1) 输出答案．

可以发现，每个节点的答案由其子树和其本身得到

我们可以先预处理出每个节点子树的大小和它的重儿子，重儿子同树链剖分一样，是拥有节点最多子树的儿子，这个过程显然可以 𝑂(𝑛) 完成．

我们用 𝑐𝑛𝑡𝑖 表示颜色 𝑖 的出现次数，𝑎𝑛𝑠𝑢 表示结点 𝑢 的答案．

遍历一个节点 𝑢，我们按以下的步骤进行遍历：

先遍历 𝑢 的轻（非重）儿子，并计算答案，但 不保留遍历后它对 𝑐𝑛𝑡 数组的影响；
遍历它的重儿子，保留它对 𝑐𝑛𝑡 数组的影响；
再次遍历 𝑢 的轻儿子的子树结点，加入这些结点的贡献，以得到 𝑢 的答案．

这样，原本全部遍历两遍，现在轻2遍，重只需要1遍，多层缩减下，整体遍历数量接近Onlogn

复制代码

#include <bits/stdc++.h>
using namespace std;

const int N = 200005;

vector<int> g[N];     // 邻接表
int c[N];             // c[u] 表示节点 u 的颜色
int sz[N];            // 子树大小
int big[N];           // 重儿子
int cnt[N];           // 全局桶：颜色出现次数
int diff;             // 当前子树内不同颜色的数量
int ans[N];           // 最终答案

/**
 * 第一遍 DFS：预处理子树大小和重儿子
 */
void dfs1(int u, int fa) {
    sz[u] = 1;
    for (int v : g[u]) {
        if (v == fa) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if (sz[v] > sz[big[u]])
            big[u] = v;
    }
}

/**
 * 子树批量更新
 * delta = +1 表示加入，-1 表示删除
 */
void apply(int u, int fa, int delta) {
    int col = c[u];

    // 维护不同颜色数
    if (cnt[col] == 0 && delta == 1) diff++;
    if (cnt[col] == 1 && delta == -1) diff--;

    cnt[col] += delta;

    // 递归处理整棵子树
    for (int v : g[u]) {
        if (v == fa) continue;
        apply(v, u, delta);
    }
}

/**
 * 第二遍 DFS：DSU on Tree 主过程
 * keep = true  表示保留当前子树信息
 * keep = false 表示处理完后清空
 */
void dfs2(int u, int fa, bool keep) {
    // 1. 递归处理所有轻儿子（不保留）
    for (int v : g[u]) {
        if (v == fa || v == big[u]) continue;
        dfs2(v, u, false);
    }

    // 2. 处理重儿子（保留）
    if (big[u])
        dfs2(big[u], u, true);

    // 3. 合并：把轻儿子子树重新加入
    for (int v : g[u]) {
        if (v == fa || v == big[u]) continue;
        apply(v, u, 1);
    }

    // 4. 加入当前节点
    apply(u, fa, 1);

    // 5. 记录答案
    ans[u] = diff;

    // 6. 如果是轻儿子，整体撤销
    if (!keep)
        apply(u, fa, -1);
}

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

    int n;
    cin >> n;

    // 读入颜色
    for (int i = 1; i <= n; i++)
        cin >> c[i];

    // 读入树边
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    // DSU on Tree
    dfs1(1, 0);
    dfs2(1, 0, true);

    // 输出答案
    for (int i = 1; i <= n; i++)
        cout << ans[i] << " \n"[i == n];

    return 0;
}