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;

}