HDU6035 2017多校第一场1003 树形DP

参考：Bahuia的博客

重点在于理清子树的层次关系，对于单个点进行逐子树更新时的转移

#pragma comment(linker, "/STACK:1024000000,1024000000") 
#include<bits/stdc++.h>
#include<stdio.h>
#include<algorithm>
#include<queue>
#include<string.h>
#include<iostream>
#include<math.h>
#include<set>
#include<map>
#include<vector>
#include<iomanip>
using namespace std;
#define ll long long
#define pb push_back
#define FOR(a) for(int i=1;i<=a;i++)

const int maxn=2e5+7;

int vis[maxn];
int color[maxn];
int siz[maxn];
ll sum[maxn];	//遍历到当前位置,颜色为x的最高一批节点为根的子树大小

vector<int>G[maxn];
ll ans;

ll dfs(int u,int fa){
	siz[u]=1;
	ll allson=0;
	int sz=G[u].size();
	for(int i=0;i<sz;i++){
		int v=G[u][i];if(v==fa)continue;
		ll pre=sum[color[u]];	//递归前的sum值
		siz[u]+=dfs(v,u);
		ll add=sum[color[u]]-pre;//新链上的块大小
		ans+=(siz[v]-add)*(siz[v]-add-1)/2;
		allson+=siz[v]-add;		//记录儿子节点v组成的连通块大小
	}
	sum[color[u]]+=allson+1;
	return siz[u];
}

int main(){
	int n,kase=0;
	while(~scanf("%d",&n)){
		memset(vis,0,sizeof vis);
		memset(sum,0,sizeof sum);
		int cnt=0;
		for(int i=1;i<=n;i++){
			scanf("%d",&color[i]);
			if(!vis[color[i]])cnt++;
			vis[color[i]]=1;
			G[i].clear();
		}
		for(int i=1;i<n;i++){
			int u,v;
			scanf("%d%d",&u,&v);
			G[u].pb(v);G[v].pb(u);
		}
		printf("Case #%d: ",++kase);
		if(cnt==1){
			printf("%lld\n",1ll*n*(n-1)/2);continue;
		}
		ans=0;
		dfs(1,-1);
		for(int i=1;i<=n;i++){
			if(!vis[i])continue;
			ans+=(n-sum[i])*(n-sum[i]-1ll)/2ll;
		}
		printf("%lld\n",1ll*n*(n-1)/2ll*cnt-ans);
	}
}