hdu 6035 Colorful Tree

题目地址
题意：定义每条边的价值为边上不同颜色的个数，求整棵树的价值。
题解：这里写链接内容
代码：

#include<iostream>
#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<vector>
#include<queue>
#include<set>
#include<algorithm>
#include<map>
using namespace std;
typedef long long int ll;
typedef pair<int,int>pa;
const int N=1e6+10;
const int mod=1e9+7;
const ll INF=1e18;
const int inf=1e4;
int read()
{
    int x=0;
    char ch = getchar();
    while('0'>ch||ch>'9')ch=getchar();
    while('0'<=ch&&ch<='9')
    {
        x=(x<<3)+(x<<1)+ch-'0';
        ch=getchar();
    }
    return x;
}
/***********************************************************/
struct node
{
    int to,next;
}edge[N<<2];
int head[N<<2];
int n,cas,num,cnt;
ll ans;
ll sum[N],size[N];
int vis[N],col[N];
void init()
{
        num=0;
        ans=0;
        cnt=0;
        memset(head,-1,sizeof(head));
        memset(vis,0,sizeof(vis));
        memset(sum,0,sizeof(sum));
}
void add(int f,int to)
{
    edge[cnt].to=to;
    edge[cnt].next=head[f];
    head[f]=cnt++;
}
void dfs(int x,int y)
{
    size[x]=1;
    sum[col[x]]++;
    ll pre=sum[col[x]];
    for(int i=head[x];i!=-1;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v==y) continue;
        dfs(v,x);
        size[x]+=size[v];
        ll conut=size[v]-(sum[col[x]]-pre);
        ans+=conut*(conut-1)/2;
        sum[col[x]]+=conut;
        pre=sum[col[x]];
    }
}
int main()
{
    cas=1;
    while(~scanf("%d",&n))
    {
        init();
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&col[i]);
            if(vis[col[i]]==0)
            {
                num++;
                vis[col[i]]=1;
            }
        }
        for(int i=1;i<n;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            add(u,v);
            add(v,u);
        }
        dfs(1,0);
        ll op=1ll*num*((ll)n*(n-1)/2);
        for(int i=1;i<=n;i++)
        {
            if(vis[i])
            {
                ll ct=n-sum[i];
                ans+=ct*(ct-1)/2;
            }
        }
            printf("Case #%d: %lld\n", cas++, op-ans);
    }
    return 0;
}