HDU 6035(树形dp)

题意略。

思路：有n * (n - 1) / 2这么多边，要枚举是不可能的，感觉和数据结构也沾不上边。再加上树上染色，以一条边上不同颜色作为这个边的值，这看起来像是算贡献那种题，和17icpc沈阳的某题有点像。

那么，我们要枚举每个边不行，不如就假设所有边上都有全部颜色，然后再减去sum(wi)[wi是第i种颜色在多少边上没出现]。

为了计算wi，我们需要知道，如果在树上的一个联通块内，不包含颜色i,那么这个联通块（假设它的点数是k）内的这k * (k - 1) / 2条边都会对wi产生贡献。

那么我们的目标就是找出所有对wi产生贡献的联通块，此时我们可以利用树的性质：siz（某个联通块） =  siz（根） -  siz（子树）；

当然，我们算贡献要以子树为单位计算贡献。

定义：sum[color[cur]] = 搜索到当前节点，颜色为color[cur]的各个最高子树的节点之和。

详见代码：

#include<bits/stdc++.h>
#define maxn 200050
using namespace std;
typedef long long LL;

LL siz[maxn],color[maxn],sum[maxn],visit[maxn],colors;
LL contri;
vector<int> graph[maxn];

LL dfs(int cur,int fa){
    siz[cur] = 1;
    LL allgap = 0;
    for(int i = 0;i < graph[cur].size();++i){
        int chi = graph[cur][i];
        if(chi == fa) continue;
        LL keep = sum[color[cur]];
        siz[cur] += dfs(chi,cur);
        LL gap = siz[chi] - (sum[color[cur]] - keep);
        contri += gap * (gap - 1) / 2;
        allgap += gap;
    }
    sum[color[cur]] += (allgap + 1);
    return siz[cur];
}
void init(){
    for(int i = 0;i < maxn;++i) graph[i].clear();
    memset(visit,0,sizeof(visit));
    memset(sum,0,sizeof(sum));
    colors = contri = 0;
}

int main(){
    LL cas = 1,n;
    while(scanf("%lld",&n) == 1){
        init();
        for(int i = 1;i <= n;++i){
            scanf("%lld",color + i);
            if(visit[color[i]] == 0){
                visit[color[i]] = 1;
                ++colors;
            }
        }
        int x,y;
        LL ans = (n * (n - 1)) / 2 * colors;
        for(int i = 0;i < n - 1;++i){
            scanf("%d%d",&x,&y);
            graph[x].push_back(y);
            graph[y].push_back(x);
        }
        dfs(1,1);
        for(int i = 1;i <= n;++i){
            if(visit[i]){
                contri += (n - sum[i]) * (n - sum[i] - 1) / 2;
            }
        }
        printf("Case #%lld: %lld
",cas++,ans - contri);
    }
    return 0;
}

其中求联通块的点数用到了取差值的方法，gap就是联通块的大小，最后的sum[color[cur]] += (allgap + 1)的意思是，这些联通块本身也是子树的一部分，在算最高子树的时候也应计入在内，那个1表示的是

根节点。

最后contri += (n - sum[i]) * (n - sum[i] - 1) / 2;这是因为我们在算联通块的时候其实都是在算子树的联通块，未能把根节点算在内，这里要补上，是因为已经到整棵树的根节点了，不会再有父节点了。

带给我的收获：

1.正面不行的时候可以考虑反面。

2.联通块与边的关系。

3.联通块size的计算。