HDU 6035 Colorful Tree

Colorful Tree
Time Limit: 6000/3000 MS (Java/Others) Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 1872 Accepted Submission(s): 787

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

计算颜色x没参与的路径数目，对于每个不包含颜色x的连通块中任意两点路径都是x不参与的贡献，对于任意一种颜色x，需要求出每个不包含x的连通块大小。
size[v]-变化值就是所求连通块大小。
具体看代码注释模拟。

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 2e5 + 10;
typedef long long LL;

LL color[MAXN], sz[MAXN], sum[MAXN], vis[MAXN];
vector <int> tree[MAXN];
LL ans;
///color[u] u节点的颜色
LL dfs(int u, int pa)
{
    sz[u] = 1;///以u为根节点的子树的大小
    LL allson = 0;///所求连通块大小总和
    int cnt = tree[u].size();///树的分支
    for (int i = 0; i < cnt; i++)
    {
        int v = tree[u][i];///u的子节点
        if (v == pa) continue;
        LL last = sum[color[u]];/// 保存递归之前的sum值

        sz[u] += dfs(v, u);///递归      以u为根节点的子树的大小
        ///sum[] 的变化值
        LL add = sum[color[u]] - last;

        /// add就是结点v为根的子树中颜色为color[u]且高度最高的若干子树的大小

        ///size-变化值  对于结点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自己
    printf("%d    %lld           %lld    %lld     %lld\n",u,sum[color[u]],ans,sz[u],allson);
    return sz[u];    ///以u为根节点的子树的大小
}


int main()
{
    int n, cs = 0;
    while (scanf("%d", &n) !=EOF)
    {
        memset(vis, 0, sizeof(vis));
        memset(sum, 0, sizeof(sum));
        int cnt = 0;
        for (int i = 1; i <= n; i++)
        {
            scanf("%I64d", &color[i]);
            ///记录有多少种颜色
            if (!vis[color[i]]) ++cnt;
            vis[color[i]] = 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);
        printf("u sum[color[u]] ans sz[u] allson\n");
        /// 只有一种颜色的特殊情况
        if (cnt == 1)
        {
            printf("%I64d\n", (LL)n * (n - 1LL) / 2LL);
            continue;
        }
        ans = 0;
        dfs(1, -1);
        /// 注意最后要对整棵树来补充所有颜色剩下的联通块(包含根结点的联通块)
        for (int i = 1; i <= n; i++)
        {
            if (!vis[i]) continue;
            ans += (n - sum[i]) * (n - sum[i] - 1LL) / 2LL;  
        }

        printf("%I64d\n", (LL)n * (n - 1LL) / 2LL * cnt - ans);
    }
    return 0;
}