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 1000000007int 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;}