【HDU6035】Colorful Tree(dfs,树形dp)

记录一个菜逼的成长。。

2017 Multi-University Training Contest - Team 1

题目链接

题目大意:

有一个n个结点的树，每个结点有一个颜色值ci。每条路径的值为这条路径上不同颜色的数量。
求总共n(n−1))2路径的权值和。

笔记

考虑每个颜色的贡献。
考虑总的答案减去非法值
总的答案ans=n∗n∗(n−1)/2 (假设每条边都经过n种颜色
sum[i]:=表示层数比当前点小的距离最近的颜色为i的节点数
遍历到当前点，往下遍历直到某一子树不包含与当前点颜色相同的点，在这个子树任意选两个点，都不会有当前点颜色的贡献，所以总答案减去子树中的路径数。

#include <bits/stdc++.h>using namespace std;#define fin freopen("D://in.txt","r",stdin)#define fout freopen("D://out.txt","w",stdout)#define rep(i,l,r) for( int i = l; i <= r; i++ )#define rep0(i,l,r) for( int i = l; i < r; i++ )#define ALL(v) (v).begin(),(v).end()#define cl(a,b) memset(a,b,sizeof(a))#define clr clear()#define pb push_back#define mp make_pair#define fi first#define se secondtypedef long long LL;typedef pair<int,int> PII;const int INF = 0x3f3f3f3f;const int maxn = 200000 + 10;int sum[maxn],son[maxn],c[maxn];LL ans ;vector<int>edge[maxn];void dfs(int u,int f){  son[u] = 1;  int add = 0,pre = sum[c[u]];  for( int i = 0; i < edge[u].size(); i++ ){    int v = edge[u][i];    if(v == f)continue;    dfs(v,u);    son[u] += son[v];    int tmp = son[v] - (sum[c[u]] - pre);    add += tmp;    pre = sum[c[u]];    ans -= (LL)tmp * (tmp - 1) / 2;  }  sum[c[u]] += add + 1;}int main(){  //fin;  int n,cas = 1;  while(~scanf("%d",&n)){    rep(i,1,n)scanf("%d",c+i);    rep(i,1,n-1){      int u,v;      scanf("%d%d",&u,&v);      edge[u].pb(v);      edge[v].pb(u);    }    ans = (LL)n * n * (n-1) / 2;    dfs(1,0);    rep(i,1,n){      int tmp = n - sum[i];      ans -= (LL)tmp * (tmp - 1) / 2;      edge[i].clr;sum[i] = 0;    }    printf("Case #%d: %lld\n",cas++,ans);  }  return 0;}