# HDU 6035-Colorful Tree（树形DP）

www.MyException.Cn  网友分享于：2013-08-16  浏览：0次
HDU 6035---Colorful Tree（树形DP）

Problem Description
There is a tree with

Input
The input contains multiple test cases.

For each test case, the first line contains one positive integers

Output
For each test case, output "Case #" in one line (without quotes), where

Sample Input
3
1 2 1
1 2
2 3
6
1 2 1 3 2 1
1 2
1 3
2 4
2 5
3 6

Sample Output
Case #1: 6
Case #2: 29

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
typedef long long LL;
const int N=200005;
int c[N],sum[N],sz[N];
vector<int> t[N];
LL ans;

int dfs(int u,int pa)
{
sz[u]=1;
int cn=t[u].size();
for(int i=0;i<cn;i++)
{
int v=t[u][i];
if(v==pa) continue;
int r=sum[c[u]];
sz[u]+=dfs(v,u);
int o=sum[c[u]]-r;
ans+=(LL)(sz[v]-o)*(LL)(sz[v]-o-1)/2;
sum[c[u]]+=sz[v]-o;
}
sum[c[u]]++;
return sz[u];
}

int main()
{
///cout << "Hello world!" << endl;
int n,Case=1;
while(scanf("%d",&n)!=EOF)
{
memset(sum,0,sizeof(sum));
for(int i=1;i<=n;i++){
scanf("%d",&c[i]);
t[i].clear();
}
for(int i=1;i<n;i++)
{
int u,v;  scanf("%d%d",&u,&v);
t[u].push_back(v);
t[v].push_back(u);
}
ans=0;
dfs(1,-1);
for(int i=1;i<=n;i++)
{
ans+=(LL)(n-sum[i])*(LL)(n-sum[i]-1)/2;
}
ans=(LL)n*(LL)(n-1)*(LL)n/2-ans;
printf("Case #%d: %lld\n",Case++,ans);
}
return 0;
}