HDU 6035 Colorful Tree （树形DP）

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

【题目大意】

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

【题解】

　　我们计算单个颜色的贡献，那么就是经过该颜色至少一次的路径数量，
　　我们统计的时候在每个点记录以其为开始的路径的答案和，
　　统计的时候计算了点自身，同时有重复计算的部分，最后减去n除以2即可
　　那么我们只要在每种颜色的虚树上统计即可。
　　对于子树的贡献需要区间修改，我们在dfs序的差分数组上更改，最后求前缀和即可。　　

【代码】

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <vector>
#include <list>
using namespace std;
const int M=200010,N=(M<<1)+10;
int n,x,y,pre
,st
,en
,c
;
vector<int> u
,v
;
long long ans
;
list<int> l
;
int dfn;
void dfs(int x,int fx){
int cx=c[x];
if(l[cx].empty())u[M+cx].push_back(x);
else u[l[cx].back()].push_back(x);
pre[x]=fx;
l[cx].push_back(x);
st[x]=dfn++;
for(int i=0;i<v[x].size();i++){
int y=v[x][i];
if(y==fx)continue;
dfs(y,x);
}l[cx].pop_back();
en[x]=dfn;
}
bool isson(int x,int y){return st[y]<=st[x]&&st[x]<en[y];}
void Dfs(int x,int d){
int pos=0;
if(x<=M){
ans[st[x]]+=n-d;
ans[st[x]+1]-=n-d;
}
for(int i=0;i<v[x].size();i++){
int y=v[x][i];
if(y==pre[x])continue;
int p=pos,size=en[y]-st[y];
while(p<u[x].size()&&isson(u[x][p],y)){
size-=en[u[x][p]]-st[u[x][p]];
p++;
}ans[st[y]]+=n-size-d;
ans[en[y]]-=n-size-d;
for(int j=pos;j<p;j++)Dfs(u[x][j],n-size);
pos=p;
}
}
int Cas=1;
int main(){
while(~scanf("%d",&n)){
for(int i=1;i<=n;i++)scanf("%d",&c[i]);
memset(ans,0,sizeof(ans));
for(int i=0;i<N;i++)v[i].clear(),u[i].clear(),l[i].clear();
dfn=1;
for(int i=1;i<n;i++){
scanf("%d%d",&x,&y);
v[x].push_back(y);
v[y].push_back(x);
}dfs(1,1);
for(int i=1;i<=M;i++){v[i+M].push_back(1);Dfs(i+M,0);}
for(int i=1;i<=n;i++)ans[i]+=ans[i-1];
long long Ans=0;
for(int i=1;i<=n;i++)Ans+=ans[st[i]];
printf("Case #%d: %lld\n",Cas++,(Ans-n)/2);
}return 0;
}