【题目链接】 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[N],st[N],en[N],c[N];
vector<int> u[N],v[N];
long long ans[N];
list<int> l[N];
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;
}