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[N];
LL ans;
int n,x,y,Cas=1,c[N],pre[N],lft[N],cut[N];
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;
}