hdu--6035--Colorful Tree

Colorful Tree

Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 1821    Accepted Submission(s): 762

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

3 1 2 1 1 2 2 3 6 1 2 1 3 2 1 1 2 1 3 2 4 2 5 3 6

 

Sample Output

Case #1: 6 Case #2: 29

官方题解：

单独考虑每一种颜色，答案就是对于每种颜色至少经过一次这种的路径条数之和。反过来思考只需要求有多少条路径没有经过这种颜色即可。直接做可以采用虚树的思想（不用真正建出来），对每种颜色的点按照 dfs 序列排个序，就能求出这些点把原来的树划分成的块的大小。这个过程实际上可以直接一次 dfs 求出。

我的思路：

给你一课树，树上有n个节点，每个节点都有自己的颜色，任意路径的权重为这条路所经过节点颜色的种类，求这棵树所有路径的权———>问题可以转化为求每一种颜色的贡献，即每种颜色有多少条路径经过—————>利用逆向思维，答案又可以转化为    (路径总数 * 颜色总数) - 各颜色 没参与的路径总数———>问题就转化为   对每一种颜色求出 所有没有经过它的路径总数———>很显然，对于每个不包含颜色x的连通块中任意两点路径都是x不参与的贡献，那么问题就转化为，对于任意一种颜色x，需要求出每个不包含x的连通块大小。 (感觉有点复杂，，，，，)那么如何求不包含x的联通块呢？

这里利用了树形dp，对于结点u，若u的颜色为x，那么dfs(u)的过程中，我们就想知道对于u的每个儿子v构成的子树中最高的一批颜色也为x的结点是哪些，要是知道这些结点子树的大小，只要拿子树v的大小减去这些节点子树的大小，就可以得到包括v在内的没有颜色x的连通块大小。网上看到一篇博文举的例子不错：

如图： 
 
对于结点1，我们想求出与1相邻的且不是黑色的结点组成的连通块大小，此时就要dfs结点1的两个儿子2和3，若我们处理出来结点2中，最高的一批黑色结点（4，8）的所构成子树的大小分别为2和1，那么我们拿2为根子树的大小5，减去2，减去1，就得到了结点2不包含黑色结点的连通块大小是5-2-1=2，也就是图中的2和5组成的连通块。同理对3，可以求得连通块大小是2（{3，6}）。 
计算出了子树u中连通块大小对答案的贡献之后，就要将当前最高一批的黑色结点更新为u，最高黑色节点构成的子树大小sum加上子树u中没有计算的那部分大小5（{1,2,3,5,6}）。 sum[x]表示的遍历到当前位置，颜色为x的高度最高一批结点为根的子树大小总和。

最后需要注意的是

之后在main函数时 ans加上的值是有可能某些节点在颜色i的父亲等较高位置，因此sum[i]无法统计到，这是上层的连通块数量也是颜色i无法贡献的，所以也应该减去！！ 

代码：

 C++ Code 

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75

#include  using  namespace std; const  int MAXN = 2e5 +  10; typedef  long  long LL; LL color[MAXN], sz[MAXN], sum[MAXN], vis[MAXN]; ///vis[*]表示某个颜色是否出现过，color[*]表示节点 * 的颜色； ///sum[x]表示的遍历到当前位置，颜色为x的高度最高一批结点为根的子树大小总和。 ///sz[X]表示 以X为根节点的  这个子树的size vector < int> tree[MAXN]; LL ans; LL dfs( int u,  int pa) ///u表示这次访问u节点，pa表示上次访问的节点 {     sz[u] =  1;                    ///以这个点为根的树的大小     LL allson =  0;               ///      int cnt = tree[u].size();    ///有多少个点与u相连通      for ( int i =  0; i < cnt; i++)     {          int v = tree[u][i];      ///下一个要访问的点          if(v==pa)                ///不能重复访问              continue;         LL last = sum[color[u]];             /// 保存递归之前的sum值sum[某种颜色]  最早以color[u]点为根到这个color[u]点之间树的大小；         sz[u] += dfs(v,u);                   ///对u的子节点进行访问；         LL add = sum[color[u]] - last;       /// add就是结点v为根的子树中颜色为color[u]且高度最高的若干子树的大小          /// 对于结点v来说，sz[v]-add就是v这棵子树最上端的，且不包含颜色为color[u]的连通块大小         ans += (sz[v] - add) * (sz[v] - add -  1) /  2;     /// 对这个连通块中任意两个点的路径都不包含颜色color[u]         allson += sz[v] - add;           /// allson记录下儿子结点v组成的不含颜色color[u]的连通块大小总和     }     sum[color[u]] += allson +  1;         /// sum更新，此时要加上不含color[u]连通块的大小总和以及u自己      return sz[u]; } int main() {      int n, cs =  0;      while (scanf( "%d", &n) !=EOF)     {         memset(vis,  0,  sizeof(vis)); ///标记颜色是否出现过，出现过为1，没有则为0，开始置为0         memset(sum,  0,  sizeof(sum)); ///sum[x]表示遍历到x位置，颜色为x的 高度 “最高“  一批节点为根的子树大小和          int cnt =  0; ///记录出现颜色的种类          for ( int i =  1; i <= n; i++)         {             scanf( "%I64d", &color[i]);      ///编号为i的节点的颜色为color[i];              if (!vis[color[i]]) ///如果出现新的颜色，则cnt++；                 ++cnt;             vis[color[i]] =  1;      ///vis[*]为1表示该种颜色出现过，进行标记             tree[i].clear();        ///这个节点的向量为零，即开始没有点与该点相连         }          for ( int i =  1; i < n; i++)         {              int u, v;             scanf( "%d%d", &u, &v);       ///建立节点之间的连接，双向的连接；             tree[u].push_back(v);             tree[v].push_back(u);         }         printf( "Case #%d: ", ++cs);          if (cnt ==  1)                /// 只有一种颜色的特殊情况         {             printf( "%I64d\n", (LL)n * (n -  1) /  2);              continue;         }         ans =  0;         dfs( 1, - 1);      ///（起始）开始遍历编号为1的节点，上一个节点为-1（表示不存在）          for ( int i =  1; i <= n; i++)         /// 注意最后要对整棵树来补充所有颜色剩下的联通块         { ///此步骤减去的是 在i颜色节点的 上面部分，有n-sum[i]个节点，这些节点互通不需要经过i颜色节点。              if (!vis[i])                  continue; ///ans加上的值是有可能某些节点在颜色i的父亲等较高位置，因此sum[i]无法统计到， ///这是上层的连通块数量也是颜色i无法贡献的，所以最后也应该减去！！             ans += (n - sum[i]) * (n - sum[i] -  1) /  2;         }         printf( "%I64d\n", (LL)n * (n -  1) /  2 * cnt - ans);     }      return  0; }