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[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;
}