U41492 树上数颜色【dsu on tree】

题目链接

题意：有一棵以 $1$ 为根结点的树，树上每个结点有一个颜色，多组询问：问以 $u$ 为根的子树上共有多少种不同的颜色。

---

解决该题，最暴力的做法就是对于每一棵子树都dfs一遍统计出答案，统计完清空：$O(n^2)$。但“听说”可以用dsu on tree优化：

dsu on tree [ $O(nlogn)$ ]

统计以结点 $u$ 为根结点的子树

1. 遍历结点 $u$ 的轻儿子，将轻儿子的贡献清空
2. 统计结点 $u$ 重儿子的贡献
3. 暴力统计轻儿子的贡献

#include <bits/stdc++.h>
using namespace std;
int read()
{
    int x = 0, f = 1; char c = getchar();
    while (c < '0' || c > '9') { if (c == '-') f = -f; c = getchar(); }
    while (c >= '0' && c <= '9') { x = x * 10 + c - '0'; c = getchar(); }
    return x * f;
}
const int maxN = 100005;
int n, m, c[maxN];  //c[i]: 颜色标号
struct EDGE{
    int adj, to;
    EDGE(int a = -1, int b = 0): adj(a), to(b) {}
}edge[maxN << 1];
int head[maxN], cnt;
void add_edge(int u, int v)
{
    edge[cnt] = EDGE(head[u], v);
    head[u] = cnt ++ ;
}

int siz[maxN], son[maxN];   //子树大小 //重儿子标号
void dfs(int u, int fa)
{
    siz[u] = 1;
    for(int i = head[u]; ~i; i = edge[i].adj)
    {
        int v = edge[i].to;
        if(v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[son[u]] < siz[v]) son[u] = v;
    }
}

int color[maxN], ans[maxN], tot, nowSon;
/*
 * color[i]: 颜色标号为i的颜色出现的次数
 * ans[i]: 以i节点为根节点的子树上的颜色总数
 * tot: 当前不同颜色的个数
 * nowSon: 当前的重儿子
 */
void cal(int u, int fa, int val)
{
    if(!color[c[u]]) ++ tot;
    color[c[u]] += val;
    for(int i = head[u]; ~i; i = edge[i].adj)
    {
        int v = edge[i].to;
        if(v == fa || v == nowSon) continue;
        cal(v, u, val);
    }
}
void dsu(int u, int fa, bool op)    //op==0: 轻边，需要清空
{
    for(int i = head[u]; ~i; i = edge[i].adj) //遍历轻儿子
    {
        int v = edge[i].to;
        if(v == fa || v == son[u]) continue;
        dsu(v, u, false);
    }
    if(son[u]) dsu(son[u], u, true), nowSon = son[u]; //统计重儿子的贡献
    cal(u, fa, 1), nowSon = 0; // 暴力统计轻儿子的贡献（记录重儿子nowSon来忽略重儿子的贡献）
    ans[u] = tot;   //得到u子树的答案
    if(!op) { cal(u, fa, -1); tot = 0; } // 如果是轻儿子，清除该子树的影响。
}

int main()
{
    memset(head, -1, sizeof(head));
    n = read();
    for(int i = 0; i < n - 1; ++ i )
    {
        int u, v; u = read(); v = read();
        add_edge(u, v); add_edge(v, u);
    }
    for(int i = 1; i <= n; ++i ) c[i] = read();
    dfs(1, 0);
    dsu(1, 0, true);
    m = read();
    for(int i = 0; i < m; ++ i )
    {
        int u = read();
        printf("%d\n", ans[u]);
    }
    return 0;
}
}

参考博客：https://www.cnblogs.com/TY02/p/11218999.html【实践证明，调试一遍就会了】
我的调试样例：

7
1 2
1 3
2 4
2 5
4 6
3 7
1 2 2 3 3 4 3
7
1
2
3
4
5
6
7