hdu 6035 Colorful Tree 树状dp

Problem Description

There is a tree with n nodes, each of which has a type of color represented by an integer, where the color of node i is ci.

The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it.

Calculate the sum of values of all paths on the tree that has n(n−1)2 paths in total.

 

Input

The input contains multiple test cases.

For each test case, the first line contains one positive integers n, indicating the number of node. (2≤n≤200000)

Next line contains n integers where the i-th integer represents ci, the color of node i. (1≤ci≤n)

Each of the next n−1 lines contains two positive integers x,y (1≤x,y≤n,x≠y), meaning an edge between node x and node y.

It is guaranteed that these edges form a tree.

 

Output

For each test case, output "Case #x: y" in one line (without quotes), where x indicates the case number starting from 1 and y denotes the answer of corresponding case.

 

Sample Input

31 2 11 22 361 2 1 3 2 11 21 32 42 53 6

 

Sample Output

Case #1: 6Case #2: 29
题意：给一颗树，定义树上路径u到v的价值为路径上不同颜色的节点的数量，问所有路径的总价值为多少
思路：首先把问题转化成,求每种颜色的贡献,每种颜色的贡献就是经过该种颜色的不同路径数目的和,反向思考一下,就是路径总数减去不经过该种颜色的路径数.
首先介绍一下代码中的几个数组的含义
sum【i】表示与颜色i相关的节点数量
例如：
首先节点1对应的sum在左子树上就是以4为根节点的子树的节点数+以8为根节点的子树的节点数（中间的2,5对应的就是黑色不参与的联通块）在计算完1这个节点之后还需要把计算过的联通量也加入到sum里边
size[i]代表的就是代表以i 为根节点的子树的节点数
#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];  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;                         int last_num = num[color[p]]; //原来与color【p】有关的节点数    int tol = 0;      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;//当前子树中与 color【p】有关的节点数         int tmp = siz[v] - (num[color[p]] - last_num); //当前子树中对应的联通量         ans += (LL)tmp * (tmp-1) / 2;            last_num = num[color[p]]; //更新last_num     }      num[color[p]] += siz[p] - tol;//把计算过的连通量也加入num }  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;          printf("Case #%d: %lld\n",cas++,ans);      }      return 0;  }