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

记录一个菜逼的成长。。

2017 Multi-University Training Contest - Team 1

题目链接

题目大意:

有一个$n$个结点的树，每个结点有一个颜色值$c_i$。每条路径的值为这条路径上不同颜色的数量。
求总共$\frac {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 second
typedef 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;
}