HDU6035-Colorful Tree 补集思想+树形DP

传送门

题意：

给一棵树，树上的每一个点都有一种颜色，设定每一条边的值为边上不同颜色点的个数，问树上所有边的总价值。

思路：

1、假设共有 all_color 种颜色，共有 all_edge = n * (n-1) / 2 条边，

则 总价值 = all_color - all_edge - 每种颜色不作用的边的总价值

2、对于一种颜色x，怎么计算没参与的路径数目呢，很显然，对于每个不包含颜色x的连通块中任意两点路径都是x不参与的贡献，那么问题就转化为，对于任意一种颜色x，需要求出每个不包含x的连通块大小

3、树形DP O（n）解决问题，具体看代码

#include <bits/stdc++.h>
using namespace std;
#define LL long long
#define N 200005
#define M 105
#define INF 0x3f3f3f3f
#define mod 1000000007
int color[N];
bool has[N];
struct edge
{
    int v,nxt;
}e[N<<1];
int head[N];
int tot;
int siz[N],num[N];//siz[i] 树形DP过程中,i点子树的大小；num[i] ,颜色i 已经去掉的点
LL ans;           //每种颜色不作用的边的总价值
inline void add_edge(int a,int b)
{
   e[tot].v = b;
   e[tot].nxt = head[a];
   head[a] = tot++;
}
void dfs(int p,int pre)
{
    siz[p] = 1;                      //初始化 子树大小为 1
    int last_num = num[color[p]];    //记录当前 与color[p]有关的 已经去掉的点
    int tol = 0;                     //记录子树 与color[p]有关的 累计去掉的点，防止重复减去
    for(int i=head[p];~i;i=e[i].nxt){
        int v = e[i].v;
        if(v==pre) continue;
        dfs(v,p);
        siz[p] += siz[v];            
        tol += num[color[p]] - last_num; //加上子树v 去掉的点
        int tmp = siz[v] - (num[color[p]] - last_num); //子树的大小 减去 子树中已经去掉的 与color[p]有关的 点
        ans += (LL)tmp * (tmp-1) / 2;    //加上 该不包含 color[p] 的连通块 的线路数
        //cout<<ans<<endl;
        last_num = num[color[p]];        //记录搜索完子树 v 的当前 与color[p]有关的 已经去掉的点
    }
    num[color[p]] += siz[p] - tol;       //更新 与color[p]有关的 已经去掉的点
}
int main()
{
    int n;
    int cas = 1;
    while(~scanf("%d",&n)){
        memset(has,0,sizeof(has));
        for(int i=1;i<=n;i++) scanf("%d",&color[i]),has[color[i]] = 1;
        memset(head,-1,sizeof(head));
        tot = 0;
        for(int i=0;i<n-1;i++){
            int a,b;
            scanf("%d%d",&a,&b);
            add_edge(a,b);
            add_edge(b,a);
        }
        ans = 0;
        memset(num,0,sizeof(num));
        dfs(1,-1);
        int all_color = 0;
        for(int i=1;i<=n;i++) all_color += has[i];
        for(int i=1;i<=n;i++){
            if(!has[i]) continue;
            int tmp = n - num[i];           //每一种颜色 剩下的点 也为 一个与该颜色 不想关的连通块
            ans += (LL)tmp * (tmp-1) / 2;   
        }
        ans = (LL)n * (n-1) / 2 * all_color - ans;  //总价值 = all_color - all_edge - 每种颜色不作用的边的总价值
        printf("Case #%d: %lld\n",cas++,ans);
    }
    return 0;
}