【HDU 6035 Colorful Tree】+ 树形 dp + 思维

Colorful Tree

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

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

Source
2017 Multi-University Training Contest - Team 1

题意 ： 有一个 n 个节点的树，每个节点都有颜色，任意两点的价值，等于两点之间路径上所有不同颜色的数目，求所有的价值和

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

AC代码：

#include<cstdio>
#include<vector>
#include<cstring>
using namespace std;
typedef long long LL;
const int MAX = 2e5 + 10;
int n,a[MAX],vis[MAX];
LL sum,s[MAX],m[MAX];
vector <int> v[MAX];
void dfs(int x,int f){
    s[x] = 1;
    LL ans = 0;
    for(int i = 0; i < v[x].size(); i++){
        int w = v[x][i];
        if(w == f) continue;
        LL o = m[a[x]];
        dfs(w,x);
        s[x] += s[w]; // 子树的节点数
        LL p = m[a[x]] - o; // 和 a[x] 相同的节点个数
        sum += (LL)(s[w] - p) * (s[w] - p - 1) / 2;
        ans += p; // 和 a[x] 相同的点的个数
    }
    m[a[x]] += s[x] - ans; // 子树中和 a[x] 不同的点个数
}
int main()
{
    int pl = 0;
    while(~scanf("%d",&n)){
        memset(vis,0,sizeof(vis));
        memset(m,0,sizeof(m));
        int nl = 0,x,y;
        for(int i = 1; i <= n ; i++){
            scanf("%d",&a[i]);
            v[i].clear();
            if(!vis[a[i]]) nl++,vis[a[i]] = 1;
        }
        for(int i = 1; i < n; i++){
            scanf("%d %d",&x,&y);
            v[x].push_back(y),v[y].push_back(x);
        }
        if(nl == 1) printf("Case #%d: %lld\n",++pl,(LL) n * (n - 1) / 2);
        else{
            sum = 0;
            dfs(1,-1);
            for(int i = 1; i <= n; i++)
                if(vis[i])
                    sum += (LL)(n - m[i]) * (n - m[i] - 1) / 2; // 相同点间的互相联通
            printf("Case #%d: %lld\n",++pl,(LL) n * (n - 1) / 2 * nl - sum);
        }
    }
    return 0;
}