【树形dp】hdu6035 Colorful Tree

非常棒的题解，我就不复述了：http://blog.csdn.net/Bahuia/article/details/76141574

O(n)

#include<cstdio>
#include<cstring>
using namespace std;
typedef long long ll;
ll ans;
bool vis[200010],hav[200010];
int sum[200010];
int v[400010],e,next[400010],first[200010];
void AddEdge(int U,int V){
	v[++e]=V;
	next[e]=first[U];
	first[U]=e;
}
int n,c[200010],siz[200010];
void dfs(int U){
	vis[U]=1;
	siz[U]=1;
	int allcU=0;
	for(int i=first[U];i;i=next[i]){
		if(!vis[v[i]]){
			int tmp=sum[c[U]];
			dfs(v[i]);
			int add=sum[c[U]]-tmp;
			siz[U]+=siz[v[i]];
			ans+=(ll)(siz[v[i]]-add)*(ll)(siz[v[i]]-add-1)/2ll;
			allcU+=(siz[v[i]]-add);
		}
	}
	sum[c[U]]+=(allcU+1);
}
int main(){
	int x,y,zu=0;
	while(scanf("%d",&n)!=EOF){
		++zu;
		ans=0;
		memset(sum,0,sizeof(sum));
		memset(first,0,sizeof(first));
		memset(vis,0,sizeof(vis));
		memset(hav,0,sizeof(hav));
		e=0;
		int cnt=0;
		for(int i=1;i<=n;++i){
			scanf("%d",&c[i]);
			if(!hav[c[i]]){
				hav[c[i]]=1;
				++cnt;
			}
		}
		for(int i=1;i<n;++i){
			scanf("%d%d",&x,&y);
			AddEdge(x,y);
			AddEdge(y,x);
		}
		dfs(1);
		for (int i=1;i<=n;++i){
			if(hav[i]){
				ans+=(ll)(n-sum[i])*(ll)(n-sum[i]-1)/2ll;
			}
        }
        printf("Case #%d: %I64d
",zu,(ll)n*(ll)(n-1ll)/2ll*(ll)cnt-ans);
	}
	return 0;
}