hdu6035[dfs+思维] 2017多校1

Posted 2020-10-01

/*hdu6035[dfs+思维] 2017多校1*/
//合并色块, 妙啊妙啊
#include<bits/stdc++.h>
using namespace std;
const double eps=1e-8;
const int inf=0x3f3f3f3f;
typedef long long LL;
vector<int>G[200005];
LL sum[200005];
int c[200005],son[200005],mark[200005];
int n,u,v,cnt=0,kase=1;
LL ans,res;
void dfs(int u,int fa){
    son[u]=1;
    LL y=sum[c[u]];
    LL x=sum[c[u]];
    for(int i=0;i<G[u].size();++i){
        int v=G[u][i];
        if(v!=fa){
            dfs(v,u);
            son[u]+=son[v];
            LL temp=son[v]-sum[c[u]]+x;
            res+=(temp-1LL)*temp/2LL;
            x=sum[c[u]];
        }
    }
    sum[c[u]]+=son[u]-(sum[c[u]]-y);
    //printf("Node #%d, Color:%d, Sum[%d]:%lld, Son:%d\n",u,c[u],c[u],sum[c[u]],son[u]);
}
void solve(){
    res=ans=0;
    dfs(1,-1);
    //cout<<"Res: "<<res<<endl;
    ans=cnt*1LL*n*(n-1)/2LL-res;
    for(int i=1;i<=n;i++){
        if(mark[i]){
            ans-=(n-sum[i])*(n-sum[i]-1)/2LL;
            //printf("Color:%d, Sum[%d]:%d, Ans:%d\n",i,i,sum[i],ans);
        }
    }
    printf("Case #%d: %lld\n",kase++,ans);
}
int main(){
    while(~scanf("%d",&n)){
        cnt=0;
        memset(sum,0,sizeof(sum));
        memset(son,0,sizeof(son));
        memset(mark,0,sizeof(mark));
        for(int i=0;i<=n+1;i++)
            G[i].clear();
        for(int i=1;i<=n;i++){
            scanf("%d",&c[i]);
            if(!mark[c[i]]){
                mark[c[i]]++;
                cnt++;
            }
        }
        for(int i=1;i<=n-1;i++){
            scanf("%d%d",&u,&v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        solve();
    }
    return 0;
}
/*
8
1 2 1 1 2 1 2 1
1 2
1 3
2 4
2 5
3 6
4 7
5 8
12
1 2 1 1 2 1 2 2 1 2 2 2
1 2
1 3
2 4
2 5
3 6
3 7
4 8
5 9
5 10
7 11
7 12
13
1 2 1 1 2 1 2 2 1 1 2 2 2
1 2
1 3
2 4
2 5
3 6
3 7
4 8
4 9
5 10
5 11
7 12
7 13
*/