HDU 6035 Colorful Tree（补集思想+树形DP）

【题目链接】 http://acm.hdu.edu.cn/showproblem.php?pid=6035

【题目大意】

　　给出一颗树，一条路径的价值为其上点权的种类数，求路径总价值

【题解】

单独考虑每一种颜色，答案就是对于每种颜色至少经过一次这种的路径条数之和。
　　反过来思考只需要求有多少条路径没有经过这种颜色即可。
　　直接做可以采用虚树的思想（不用真正建出来），对每种颜色的点按照 dfs 序列排个序，
　　就能求出这些点把原来的树划分成的块的大小。
　　在搜索的过程中我们保存每个颜色的父节点，一遍dfs即可得到答案。

【代码】

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <vector>
using namespace std;
typedef long long LL;
const int N=200010;
vector<int> v
;
LL ans;
int n,x,y,Cas=1,c
,pre
,lft
,cut
;
LL sum2(LL x){return x*(x-1)/2;}
int dfs(int x,int fx){
int res=1,fa=pre[c[x]];
pre[c[x]]=x;
for(int i=0;i<v[x].size();i++){
int y=v[x][i];
if(y==fx)continue;
cut[x]=0;
int sz=dfs(y,x);
res+=sz;
ans-=sum2(sz-cut[x]);
}(fa?cut[fa]:lft[c[x]])+=res;
pre[c[x]]=fa;
return res;
}
int main(){
while(~scanf("%d",&n)){
for(int i=1;i<=n;i++)scanf("%d",&c[i]),v[i].clear(),pre[i]=cut[i]=lft[i]=0;
for(int i=1;i<n;i++){
scanf("%d%d",&x,&y);
v[x].push_back(y);
v[y].push_back(x);
}ans=sum2(n)*n; dfs(1,1);
for(int i=1;i<=n;i++)ans-=sum2(n-lft[i]);
printf("Case #%d: %lld\n",Cas++,ans);
}return 0;
}