2017 Multi-University Training Contest - Team 1 1003 Colorful Tree

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=6035

题目大意： 给一棵n个点树， 每一个点有一个颜色ci， 对于每一条边对答案的贡献为这一条边上颜色不同的个数， 问有的n * (n - 1) / 2条边的贡献和是多小。

参考博客：http://blog.csdn.net/my_sunshine26/article/details/76163937

解题思路： 首先可以想如果每一个点对于每一条边都有贡献， 那么总答案为n * n * (n - 1) / 2， 但是有一些颜色对于一些边来说是没有贡献的， 那么我们就要考虑如何减去。

  我们用num[i]表示以i为“根节点”，其子树的节点个数和(包括自身，且整棵树的根节点，即1，事先已经确定)。

再用sum[color[u]]表示搜索到u点时，以其为根节点的每棵子树中距离u最近且颜色相等的节点x的num[x]值的和。

容易知道，每搜索一次u的根节点v，sum[color[u]]的值会变大temp，且num[v]-temp的值即表示这些点之间的路径不经过color[u]这种颜色，这种路径的数量为temp*(temp-1)/2,应减去，重复此操作即可。

#include<cstdio>
#include<cstring>
#include<cmath>
#include<cstdlib>
#include<vector>
#include<algorithm>
using namespace std;
const int MaxN = 2e6;
typedef long long LL;

LL ans, n;  
LL num[MaxN + 5], sum[MaxN + 5];
int color[MaxN + 5];
vector<int> edge[MaxN + 5];

void dfs(int u, int fa){
	int add = 0;
	num[u] = 1;
	for(int i = 0; i < edge[u].size(); i++){
		int v = edge[u][i];
		if(v == fa) continue;
		int pre = sum[color[u]];
		dfs(v, u);
		num[u] += num[v];
		int temp = num[v] - (sum[color[u]] - pre);
		ans -= (LL)temp *(temp - 1) / 2;
		add += temp;
	}
	sum[color[u]] += add + 1;
}

int main(){
	int Case = 0;
	while(~scanf("%lld", &n)){
		memset(num, 0, sizeof(num));
		memset(sum, 0, sizeof(sum));
		for(int i = 1; i <= n; i++) edge[i].clear();
		ans = 0;
		for(int i = 1; i <= n; i++) scanf("%d", &color[i]);
		for(int i = 1; i < n; i++){
			int u, v;
			scanf("%d %d", &u, &v);
			edge[u].push_back(v);
			edge[v].push_back(u);
		}
		ans = n * n * (n - 1LL) / 2LL;
		dfs(1, -1);
		for(int i = 1;i <= n;i++){
            int temp = n - sum[i];  
            ans -= (LL)temp * (temp - 1) / 2;  
        } 
        printf("Case #%d: %lld\n",++Case,ans);  
	}
	return 0;
}