题意略。
思路：有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的计算。