2017 多校训练第一场 HDU 6035 Colorful Tree

#include <bits/stdc++.h>
using namespace std;
#define test TEST 
typedef long long ll;
const int maxn=2e5+5;

int color[maxn],sum[maxn],siz[maxn];
vector<int>edge[maxn];
ll ans;

void dfs(int u,int fa){
	siz[u]=1;
	int i,v,t=color[u],pre=sum[color[u]];
	int add=0;
	for(i=0;i<edge[u].size();i++){
		v=edge[u][i];
		if(v==fa) continue;
		dfs(v,u);
		siz[u]+=siz[v];

		ll temp=siz[v]-(sum[color[u]]-pre);
		pre=sum[color[u]];
		add+=temp;
		ans-=temp*(temp-1)/2;
	}
	sum[color[u]]+=(add+1);
}

int main(int argc, char const *argv[])
{
	int cas=1,x,y;
	ll n;
	while(~scanf("%lld",&n)){
		memset(sum,0,sizeof(sum));
		memset(siz,0,sizeof(siz));
		edge[0].clear();
		for(ll i=1;i<=n;i++){
			scanf("%d",&color[i]);
			edge[i].clear();
		}
		for(ll i=0;i<n-1;i++){
			scanf("%d %d",&x,&y);
			edge[x].push_back(y);
			edge[y].push_back(x);
		}
		ans=n*(n-1)/2;
		ans*=n;
		dfs(1,0);
		for(ll i=1;i<=n;i++){
			ll temp=n-sum[i];
			ans-=temp*(temp-1)/2;
			edge[i].clear();
		}
		printf("Case #%d: %lld\n",cas++,ans);
	}
	return 0;
}