HDU 6035(2017多校第一场)。color tree （树形dp）

Colorful Tree Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others) Total Submission(s): 1999    Accepted Submission(s): 853 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

题意：求每种颜色参与了多少路径。。这个题很明显是求颜色的贡献，但是不好求，我们可以反过来思考，用总的，减去这个颜色没有参与的路径，不就是正确结果了么。 没有该颜色参与的路径，不就是求 不存在该颜色的连通块么。

所以这个题就转化为，对于每种颜色，求出不含其的连通块。

2e5 * 2e5答案要用longlong

另外贴一个讲解比较好的博客 http://blog.csdn.net/Bahuia/article/details/76141574

#include <iostream>
#include <algorithm>
#include <string.h>
#include <stdio.h>
#include <cmath>
#include <queue>
#include <utility>
using namespace std;
using ll = long long;
const int maxn = 200005;
int n;
int color[maxn];
//sum 代表当前结点 的子树中 跟自己颜色相同的最高 子节点 which 作为根的子树大小的总和
ll sum[maxn];
ll size[maxn];
vector<int> g[maxn];
int visited[maxn];
ll ans;
ll dfs(int u, int fa) {
    size[u] = 1;//u结点本身 + 其儿子的 数量
    ll allson = 0;//u结点的所有儿子
    for (int i = 0; i < g[u].size(); ++i) {
        int v = g[u][i];
        if (v == fa)
            continue;
        ll pre_sum = sum[color[u]];//递归前的sum值
        size[u] += dfs(v, u);
        ll add = sum[color[u]] - pre_sum;//add是当前结点的 最高跟自己颜色相同子节点 which 作为根节点的树 的大小
        ans += ((size[v] - add) * (size[v] - add - 1)) / 2;//为啥是size[v]，因为求非color[u]的连通块，所以用的当前结点的子节点的size表示
        allson += size[v] - add;
    }
    sum[color[u]] += allson + 1;//+1是包括自己，现在的sum[color[u]]是已经访问过的color[u]作为根节点的树的大小（虚树的大小
    return size[u];
}
int main() {
        //freopen("in.txt", "r", stdin);
        //cout << static_cast<long long>(INT_MAX) / 2 * 100 << endl;
    int ca = 1;
    while (~scanf("%d", &n)) {
        memset(sum, 0, sizeof(sum));
        memset(size, 0, sizeof(size));
        memset(visited, 0, sizeof(visited));
        int cnt = 0;
        
        for (int i = 1; i <= n; ++i) {
                //cin >> color[i];
            scanf("%d", &color[i]);
            if (!visited[color[i]]) {
                visited[color[i]] = 1;
                cnt++;
            }
            g[i].clear();
        }
        for (int i = 1; i < n; ++i) {
            int x, y;
                //cin >> x >> y;
            scanf("%d%d", &x, &y);
            g[x].push_back(y);
            g[y].push_back(x);
        }
        printf("Case #%d: ", ca++);
        if (cnt == 1) {
            printf("%lld\n", (static_cast<ll>(n) * (n - 1)) / 2);
        } else {
            ans = 0;
            dfs(1, -1);
            for (int i = 1; i <= n; ++i) {
                if (visited[i]) {
                    ans += (n - sum[i]) * (n - sum[i] - 1) / 2;
                }
            }
            printf("%lld\n", (static_cast<ll>(n) * (n-1) / 2 * cnt - ans));
        }
    }
    
}