【dp】树形dp真好玩，hdu6035多校第一场的 colorful tree

本文链接：https://blog.csdn.net/cww97/article/details/76376831

hdu6035

贴两个很清真的blog

wy_2016

bahuia

这个故事教导我们写blog要带图

两个blog都提到了不包含该颜色的联通块，

单个联通块里的计数很容易理解，

最后dfs完了，sum还是一通狂减，

乍一看很玄学，其实是包含根结点的联通块，现场的时候就卡在这一点上

#include <vector>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long LL;
const int N = 2e5 + 801;
LL n, ans;
LL sum[N];
int col[N], size[N];
bool vis[N];
vector<int>G[N];

inline LL Cn2(LL n){ //组合数C(n,2) 
    return n*(n-1) >> 1;
}

void dfs(int u, int fa){
    int all = 0;
    size[u] = 1;
    for (int i = 0; i < G[u].size(); i++){
        int to = G[u][i];
        if (to == fa) continue;
        LL last = sum[col[u]];
        dfs(to, u);
        size[u] += size[to];
        LL part = sum[col[u]] - last;
        all += part;
        ans -= Cn2(size [to] - part);
    } 
    sum[col[u]] += size[u] - all;
}


int main(){
    //freopen("in.txt", "r", stdin);
    for (int x, y, _ = 1; ~scanf("%lld", &n); _++){
        memset(vis, 0, sizeof(vis));
        int col_cnt = 0;
        for (int i = 1; i <= n; i++){
            scanf("%d", &col[i]);
            if (!vis[col[i]]) {
                vis[col[i]] = true;
                col_cnt++;
            }
        }
        for (int i = 1; i <= n; i++) G[i].clear();
        for (int i = 1; i < n; i++){
            scanf("%d%d", &x, &y);
            G[x].push_back(y);
            G[y].push_back(x);
        }

        ans = col_cnt * Cn2(n);
        memset(sum, 0, sizeof(sum));
        dfs(1,0);
        for (int i = 1; i <= n; i++) if(vis[i]){
            ans -= Cn2(n - sum[i]);
        }
        printf("Case #%d: %lld\n", _, ans);
    }
    return 0;
}