洛谷U41492 树上数颜色树上启发式合并

野性酷女 2023-02-20 03:12 29阅读 0赞

题目链接：https://www.luogu.com.cn/problem/U41492  
题意：给一棵根为1的树，每次询问子树颜色种类数  
思路：树上启发式合并，每次计算时如果是第一次出现就计算贡献，轻链的话要消除贡献

#include<bits/stdc++.h>
    using namespace std;
    typedef long long ll;
    #define fi first
    #define se second
    #define ls rt << 1
    #define rs rt << 1|1
    #define po pop_back
    #define pb push_back
    #define mk make_pair
    #define lson l, mid, ls
    #define rson mid + 1, r, rs
    #define pll pair<ll, ll>
    #define pii pair<int, int>
    #define ull unsigned long long
    #define pdd pair<double, double>
    const int mod = 1e9 + 9;
    const int maxn = 2e5 + 10;
    const int inf = 0x3f3f3f3f;
    const ll linf = 0x3f3f3f3f3f3f3f3f;
    int c[maxn], num[maxn], sz[maxn], son[maxn], sum[maxn], cnt[maxn], ans[maxn], Max, Son;
    vector<int> e[maxn];
    
    void dfs(int u, int fa) //统计重儿子信息
    {
        sz[u] = 1;
        //sum[u] = num[u];
        for(auto v : e[u])
        {
            if(v == fa) continue;
            dfs(v, u);
            sz[u] += sz[v];
            //sum[u] += sum[v];
            if(sz[son[u]] < sz[v]) son[u] = v;
        }
    }
    void calc(int u, int fa, int val)
    {
        if(cnt[c[u]] == 0) ++Max;
        cnt[c[u]] += val;
        for(auto v : e[u])
        {
            if(v == fa || v == Son) continue;
            calc(v, u, val);
        }
    }
    void dfs2(int u, int fa, int op) // 0轻链 1重链
    {
        for(auto v : e[u])
        {
            if(v == fa || v == son[u]) continue;
            dfs2(v, u, 0);
        }
        if(son[u]) //重儿子计算贡献 不消除
             dfs2(son[u], u, 1), Son = son[u];
        calc(u, fa, 1), Son = 0;
        ans[u] = Max;
        if(!op) calc(u, fa, -1), Max = 0; //消除贡献
    }
    
    int main()
    {
        int n;
        scanf("%d", &n);
        for(int i = 1; i < n; ++i)
        {
            int u, v;
            scanf("%d%d", &u, &v);
            e[u].pb(v), e[v].pb(u);
        }
        for(int i = 1; i <= n; ++i) scanf("%d", &c[i]);
        dfs(1, 0);
        dfs2(1, 0, 0);
        int m;
        scanf("%d", &m);
        while(m--)
        {
            int x;
            scanf("%d", &x);
            printf("%d\n", ans[x]);
        }
    	return 0;
    }