HDU 6035---Colorful Tree（树形DP）

#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;
}