HDU 6035 Colorful Tree （2017多校1 - dfs思路）

题意：

给你一棵树， 树上的路权为  这条路上不同颜色的个数。  求所有路径的权值之和。

思路：

感觉好难。。

其实这个题一遍dfs 就可以办了。

可以反过来思考，对于每一种颜色求出 有多少条路 没有经过这个颜色。 最后用总路数减去即可。

令sum[i] 表示 以i 颜色为根的子树总大小。

siz[u]表示以u 为根的子树的大小。

怎么找呢？

假如我们枚举到u结点了。 对于u 的每个孩子， 我们找出他们的sum[col[u] ]   中间的位置差就是一个连通分块， 他们组成的路 都不包含颜色u。

这样把所有孩子都累加起来， 在更新sum[col[u] ]。

但是这样统计 会算出整体sum[col[u] ] 并不是这条链的。

这里的pre 变量就比较巧妙了。

在dfs 这个点之前，我们令pre = sum[ col[u] ] 就是  这个点之前的所有sum[col[u]]

然后更新完u的孩子后的sum[col[u] ]  - pre  就是这条链的。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;


const int maxn = 200000 + 10;
vector<int>g[maxn];
int col[maxn];
long long sum[maxn];
int siz[maxn];
int vis[maxn];
int num;
long long ans;
void dfs(int cur, int fa){
    siz[cur] = 1;
    int pre = sum[col[cur] ];
    bool ok = 0;
    for (int i = 0; i < g[cur].size(); ++i){
        int v = g[cur][i];
        if (v == fa) continue;
        ok = 1;
        dfs(v, cur);
        siz[cur] += siz[v];
        long long tmp = siz[v] - (sum[col[cur] ] - pre);
        ans += (long long)tmp*(long long)(tmp-1) / 2;
        sum[col[cur] ] += tmp;
        pre = sum[col[cur] ];
    }
    sum[col[cur] ] += 1;
}


int main(){
    int n, ks = 0;
    while(~scanf("%d",&n)){
        memset(vis,0,sizeof vis);
        num = ans = 0;

        for (int i = 1; i <= n; ++i){
            scanf("%d", &col[i]);
            if (!vis[col[i] ]){
                vis[col[i] ] = 1;
                ++num;
            }
            sum[i] = 0;
            g[i].clear();
        }



        for (int i = 1; i < n; ++i){
            int x,y;
            scanf("%d %d",&x, &y);
            g[x].push_back(y);
            g[y].push_back(x);
        }
        ans = 0;
        dfs(1, -1);
        for (int i = 1; i <= n; ++i){
            if (vis[i]){
                long long tmp = siz[1] - sum[i];
                ans += tmp * (tmp-1) / 2;
            }
        }

        ans = (long long)num * (long long)n * (long long)(n-1) / 2 - ans;

        printf("Case #%d: %I64d\n", ++ks, ans);
    }


    return 0;
}