洛谷U41492 树上数颜色【dsu on tree】

题意

给一棵根为1的树，每次询问子树颜色种类数

输入格式

第一行一个整数n，表示树的结点数

接下来n-1行，每行一条边

接下来一行n个数，表示每个结点的颜色c[i]

接下来一个数m，表示询问数

接下来m行表示询问的子树

输出格式

对于每个询问，输出该子树颜色数

样例 #1

样例输入 #1

5
1 2
1 3
2 4
2 5
1 2 2 3 3
5
1
2
3
4
5

样例输出 #1

3
2
1
1
1

思路

树上启发式合并模板题。
第一次学习树上启发式合并，参考了这位大佬的资料，个人认为讲的很清晰了，有手模过程好评。个人对树上启发式合并的初步理解（由于是第一道题所以只有初步理解）就是，与子树查询有关，且查询不带修改，然后维护全局信息（对于这道题是出现的颜色数量与每个颜色的出现次数），用于计算子树答案。
对于以$u$为根的子树答案的计算，首先递归求取$u$所有轻儿子的答案，然后清除所有轻儿子答案求取过程中对全局信息的影响，求取重儿子的答案，再求取$u$的答案（需重新递归轻儿子计算，重儿子就不需要了因为该影响没有清除）。
对于本题，维护一个全局的每个颜色出现次数数组$cnt$，维护出现的颜色个数$top$，按照板子来即可。

代码

#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdio>
#include <set>
#include <vector>
#include <map>
#include <queue>
#include <cmath>

using namespace std;

typedef long long LL;
typedef pair<int, int> PII;

const int N = 100005;

int n, q;
int c[N], ans[N];
int siz[N], son[N], cnt[N], top;

struct edge
{
    int to, next;
}e[N << 1];
int head[N], tot;
bool vis[N];

void add(int u, int v)
{
    e[++ tot] = {v, head[u]};
    head[u] = tot;
}

void dfs(int u, int ff)
{
    siz[u] = 1;
    for (int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].to;
        if (v == ff) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if (!son[u] || siz[v] > siz[son[u]]) son[u] = v;
    }
}

void calc(int u, int ff, int k)
{
    if (k > 0)
    {
        if (!cnt[c[u]]) top ++;
        cnt[c[u]] ++;
    }
    else 
    {
        if (cnt[c[u]] <= 1) top --;
        cnt[c[u]] --;
    }
    for (int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].to;
        if (v == ff || vis[v]) continue;
        calc(v, u, k);
    }
}

void dfs2(int u, int ff, int keep)
{
    for (int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].to;
        if (v == ff || v == son[u]) continue;
        dfs2(v, u, 0);
    }
    if (son[u]) dfs2(son[u], u, 1), vis[son[u]] = true;
    calc(u, ff, 1);
    vis[son[u]] = false;
    ans[u] = top;
    if (!keep) calc(u, ff, -1);
}

int main()
{
    scanf("%d", &n);
    for (int i = 1; i < n; i ++ )
    {
        int u, v;
        scanf("%d%d", &u, &v);
        add(u, v);
        add(v, u);
    }
    for (int i = 1; i <= n; i ++ ) scanf("%d", &c[i]);
    dfs(1, 0);
    dfs2(1, 0, 0);
    scanf("%d", &q);
    while (q --)
    {
        int x;
        scanf("%d", &x);
        printf("%d\n", ans[x]);
    }
    return 0;
}