hdu6035[dfs+思维] 2017多校1

/*hdu6035[dfs+思维] 2017多校1*/
//合并色块, 妙啊妙啊
#include<bits/stdc++.h>
using namespace std;
const double eps=1e-;
const int inf=0x3f3f3f3f;
typedef long long LL;
vector<int>G[];
LL sum[];
int c[],son[],mark[];
int n,u,v,cnt=,kase=;
LL ans,res;
void dfs(int u,int fa){
son[u]=;
LL y=sum[c[u]];
LL x=sum[c[u]];
for(int i=;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=;
dfs(,-);
//cout<<"Res: "<<res<<endl;
ans=cnt*1LL*n*(n-)/2LL-res;
for(int i=;i<=n;i++){
if(mark[i]){
ans-=(n-sum[i])*(n-sum[i]-)/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=;
memset(sum,,sizeof(sum));
memset(son,,sizeof(son));
memset(mark,,sizeof(mark));
for(int i=;i<=n+;i++)
G[i].clear();
for(int i=;i<=n;i++){
scanf("%d",&c[i]);
if(!mark[c[i]]){
mark[c[i]]++;
cnt++;
}
}
for(int i=;i<=n-;i++){
scanf("%d%d",&u,&v);
G[u].push_back(v);
G[v].push_back(u);
}
solve();
}
return ;
}
/*
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
*/
上一篇:【hihoCoder】1041. 国庆出游


下一篇:D. Eternal Victory(dfs + 思维)