U41492 dsu on tree

题意

传送门 U41492 树上数颜色

题解

预处理各子树的答案，朴素的方法即对每个节点向下 $DFS$ 一次，时间复杂度 $O(N^2)$。考虑到节点的答案为子树信息的叠加，使用树上启发式合并算法。

基本思想是递归求解儿子的答案，将重儿子（即节点数最多的儿子）放到最后求解，且不清空它所记录的信息，此时可以保留的信息是最多的；那么求解当前节点的答案时，只需要再统计一次非重儿子的信息。

定义连向重儿子的为重边，其余为轻边。由于轻边连接的子树节点个数不超过父节点所在子树的节点个数的一半，则任一节点至根节点的轻边数不超过为 $\log N$；对于任一节点，除了求解其答案需要遍历一次，还需要遍历它到根节点的轻边数次（考虑清空记录信息，对于每条轻边还要多遍历一次节点），故时间复杂度为 $O(N\log N)$。

#include <bits/stdc++.h>
using namespace std;
const int maxn = 100005;
int N, M, sum, cnt[maxn], col[maxn], sz[maxn], hs[maxn], res[maxn];
int tot, head[maxn], to[maxn << 1], nxt[maxn << 1];

inline int read()
{
    int x = 0, f = 1;
    char c = 0;
    for (; c < '0' || c > '9'; c = getchar())
        if (c == '-')
            f = -f;
    for (; c >= '0' && c <= '9'; c = getchar())
        x = (x << 1) + (x << 3) + c - '0';
    return x * f;
}

inline void add(int x, int y) { to[++tot] = y, nxt[tot] = head[x], head[x] = tot; }

void pdfs(int x, int f)
{
    sz[x] = 1;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f)
        {
            pdfs(y, x);
            sz[x] += sz[y];
            if (sz[y] > sz[hs[x]])
                hs[x] = y;
        }
    }
}

void add(int x, int f, int s, int d)
{
    int t = (cnt[col[x]] += d);
    if (t == 0)
        --sum;
    if (t == 1 && d == 1)
        ++sum;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && y != s)
            add(y, x, s, d);
    }
}

void dfs(int x, int f, int keep)
{
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && y != hs[x])
            dfs(y, x, 0);
    }
    if (hs[x])
        dfs(hs[x], x, 1);
    add(x, f, hs[x], 1);
    res[x] = sum;
    if (!keep)
        add(x, f, 0, -1);
}

int main()
{
    scanf("%d", &N);
    for (int i = 1, x, y; i < N; ++i)
        x = read(), y = read(), add(x, y), add(y, x);
    for (int i = 1; i <= N; ++i)
        col[i] = read();
    pdfs(1, 0);
    dfs(1, 0, 1);
    scanf("%d", &M);
    for (int i = 1; i <= M; ++i)
        printf("%d\n", res[read()]);
    return 0;
}