2017 Multi-University Training Contest - Team 1 ：Colorful Tree

Colorful Tree

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

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
  
  

题目意思：有一个节点个数为n的树， n个节点都有自己的颜色，从节点i到节点j路径

的权值是i到j路程中经过点的不同颜色数，求图中所有路径的权值之和。

思路，官方题解和网上的都是下面的是说法：

可以把问题转化为对每种颜色有多少条不同的路径至少经过这种颜色的点，然后加和。

求有多少条路径经过可以转换为总路径数-没有经过的路径数，只要求出没有经过的路径数就好了。

这个题目真的很难，昨天看了一天，越看越迷糊，对于网上的代码：就是不明白他们的sum数组的变化

过程，以及presum的变化过程，感觉没有人说得清楚sum数组的具体含义，看了一天后明白了，但是

现在心里明白，还是不能很清楚的解释出sum数组的含义。下面图的讲解是别人的，昨天也是看了这个

才有所醒悟。

对于结点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}）。

原文链接：http://blog.csdn.net/Bahuia/article/details/76141574

AC代码：

#include <iostream>
#include <stdio.h>
#include <vector>
#include <string.h>
using namespace std;

typedef long long LL;
const int maxn = 200010;
int color[maxn];   ///存放每个节点对应的颜色
int sz[maxn];      ///size[i]存放以i为根节点的子树的大小
int vis[maxn];     ///标记那些颜色出现过
LL sum[maxn];      ///sum[i]用来统计，当前搜索过程中被颜色i管理的点的个数
vector<int>node[maxn];  ///存放图
LL ans;
void dfs(int u,int father)
{
    sz[u] = 1;
    LL allson = 0;
    for(int i = 0; i < node[u].size(); i++)
    {
        int v = node[u][i];
        if(v == father) continue;
        LL persum = sum[color[u]];
        ///保存一下，搜索节点u之前，被color[u]这种颜色管理的点的数目
        dfs(v,u);
        sz[u] += sz[v];
        ///以节点u为根的树的大小，等于以其各个子树为根的树的大小的和再加1.
        LL ccount = sz[v]-sum[color[u]]+persum;
        ///ccount是构成连通快中的节点数目
        ans -= ccount*(ccount-1)/2;
        allson += ccount;  ///计算过后，这些点也被color[u]所管理
    }
    sum[color[u]] += allson + 1;      ///加1是加上它自己
}
int main()
{
    int n,Case=0;
    while(~scanf("%d",&n))
    {
        memset(sum,0,sizeof(sum));
        memset(sz,0,sizeof(sz));
        memset(vis,0,sizeof(vis));
        LL num = 0;
        for(int i = 1; i <= n; i++)
        {
            scanf("%d",&color[i]);
            node[i].clear();
            if(!vis[color[i]])
            {
                vis[color[i]] = 1;
                num++;
            }
        }
        for(int i = 1; i < n; i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            node[u].push_back(v);
            node[v].push_back(u);
        }
        ans = num*n*(n-1)/2;
        dfs(1,0);
        for(int i = 1; i <= n; i++)
        {
            if(vis[i])
            {
                LL cnt = n-sum[i];
                ans = ans - cnt*(cnt-1)/2;
            }
        }
        printf("Case #%d: %lld\n",++Case,ans);
    }
    return 0;
}
/*
15
4 3 1 2 2 3 5 1 7 10 4 3 2 1 5
1 2
1 3
2 4
2 5
2 6
3 7
3 8
5 9
5 10
6 11
7 12 
7 13
8 14
13 15

368
*/