提交通道
洛谷日報

考慮非\(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;
}
 

洛谷U41492（樹上啟發式合並）