洛谷U41492（树上启发式合并）

提交通道
洛谷日报

考虑非\(O(n^2)\)的预处理。一遍dfs时，check某颜色有没有的数组何时清空很尴尬：得到某树答案后如果不清，则影响接下来兄弟树的搜索；如果清了，父亲节点又难以收集答案。
解决方法：先让儿子们各顾各的家，算一遍各自的答案（假如能算），check清就清了吧。然后考虑人为优化，即重链求完后等一等！先别清！然后将轻链重新扫一遍，也不清check数组的。代码中的keep就控制是否要清。这样轻链扫两遍，重链扫一遍，就得到了儿子们和父亲的答案，随机数据下复杂度\(O(nlogn)\)。

const int maxn = 1e5 + 5;
int n, m, u, v, c[maxn];
int size[maxn], son[maxn], ans[maxn];
bool check[maxn];
vector<int> adj[maxn];

void dfs1(int cur, int fa) {
    size[cur] = 1;
    for (int i : adj[cur])
        if (i != fa) {
            dfs1(i, cur);
            size[cur] += size[i];
            if (size[i] > size[son[cur]])
                son[cur] = i;
        }
}

int dfs2(int cur, int fa, int keep) {
    if (keep) {
        for (int i : adj[cur])
            if (i != fa && i != son[cur])
                dfs2(i, cur, keep);
    }
    int res = 0;
    if (son[cur])   res += dfs2(son[cur], cur, keep);
    for (int i : adj[cur])
        if (i != fa && i != son[cur])
            res += dfs2(i, cur, 0);

    if (!check[c[cur]]) res++, check[c[cur]] = 1;
    if (keep) {
        ans[cur] = res;
        if (cur != son[fa]) mset(check, 0);
    }

    return res;
}

int main() {
    read(n);    
    rep(i, 1, n - 1) {
        read(u), read(v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    rep(i, 1, n)    read(c[i]);

    dfs1(1, 0);
    dfs2(1, 0, 1);

    for (read(m); m--; ) {
        read(u);
        writeln(ans[u]);
    }
    return 0;
}