HDU-6035 Colorful Tree（树型dp）

Colorful Tree

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

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
  
  

题目要求每条路径经过的不同颜色数之和，可以转换成每种颜色对答案的贡献
“有多少条路径经过一个或多个颜色x的点”可以转换成“有多少条路径不经过颜色x”
即：经过颜色x的路径数量=n(n-1)/2-不经过颜色x的路径数

怎么求不经过颜色x的路径数呢？
我们可以想象把颜色为x的点选出来，这些点会把树分成若干个子树，这些子树内的路径都是不会经过颜色为x的点的
这样我们通过dfs遍历到顶点u，颜色为x，在它的子树内所有以颜色x为根的子树都要舍去

这一做法很简单：维护一个sum[i]表示以颜色i为根的子树总大小，每次遍历u的子节点v时，sum[col[u]]就会增加，
其增加量tmp就是子树v中以颜色col[u]为根的子树大小，用sz[v]-tmp就是子树v中颜色部位col[u]的点

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const int MX = 2e5 + 5;
const LL mod = 1e9 + 7;

struct Edge {
    int v, nxt;
} E[MX * 2];
int head[MX], rear;
void add(int u, int v) {
    E[rear].v = v;
    E[rear].nxt = head[u];
    head[u] = rear++;
}

int col[MX];
int sz[MX];
int vis[MX];
LL sum[MX];
LL ans;

void init(int n) {
    for (int i = 1; i <= n; i++) head[i] = -1, vis[i] = sum[i] = 0;
    rear = 0;
}


void dfs(int u, int fa) {
    sz[u] = 1;
    LL pre = sum[col[u]], c = 0, szv;
    for (int i = head[u]; ~i; i = E[i].nxt) {
        int v = E[i].v;
        if (v == fa) continue;
        dfs(v, u);
        sz[u] += sz[v];
        LL tmp = sum[col[u]] - pre;//子树v中以颜色col[u]为根的子树大小
        szv = sz[v] - tmp; //子树v中其它颜色的连通的节点数量
        pre = sum[col[u]];
        ans -= szv * (szv - 1) / 2;
        c += tmp;  //c为了避免重复计算sz[u]的大小
    }
    sum[col[u]] += sz[u] - c; //sum[col[u]]加上以u为根的子树大小
}

int main() {
    int n, cas = 0;
    //freopen("in.txt", "r", stdin);
    while (~scanf("%d", &n)) {
        init(n);
        int cnt = 0;
        for (int i = 1; i <= n; i++) {
            scanf("%d", &col[i]);
            if (!vis[col[i]]) {
                vis[col[i]] = 1;
                cnt++;
            }
        }
        for (int i = 1, u, v; i < n; i++) {
            scanf("%d%d", &u, &v);
            add(u, v); add(v, u);
        }
        ans = (LL)cnt * n * (n - 1) / 2;
        dfs(1, -1);
        for (int i = 1; i <= n; i++) {
            if (sum[i] == 0) continue;
            LL tmp = n - sum[i];
            ans -= tmp * (tmp - 1) / 2;
        }
        printf("Case #%d: %lld\n", ++cas, ans);
    }
    return 0;
}