hdu6035[dfs+思维] 2017多校1

/*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
*/

posted @ 2017-08-06 21:07 UnderSilence 阅读(...) 评论(...) 编辑 收藏