Colorful Tree

HDU - 6035


给定一棵树，每个点有一个颜色1~n，定义两点的路径是两点间不同的颜色数量，求n*(n-1)/2条路径的权值总和。
一开始想主席树，但是后来发现并不用那么麻烦。
假设考虑颜色1的贡献值，可以把所有颜色为1的点割掉，用n*(n-1)/2 减去剩下的所有子树的路径和，就是所有经过了颜色1的路径数，也就是路径1的贡献。那么怎么快速求每一颗子树的大小，考虑肯定是从下往上进行分割，我用fa[t]代表当前dfs到x的节点，往上走第一个颜色为t的节点是fa[t]，然后我们可以把x节点压如到fa[t]节点所建立的vector中，当回溯到fa[t]时，减去所有vector中的子树大小，就是割剩的树的大小m，路径总数就是m*(m-1)/2，当fa[t]处理另一颗子树时vector要清空。注意就是当fa[t]是0，即是从上往下第一个颜色是t的节点时要特殊处理。

#include <bits/stdc++.h>#include <vector>using namespace std;#define ll long longconst int maxn = 200010;vector<int> G[maxn];vector<int> s[2*maxn], c[maxn];int a[maxn], n, x, y, r[maxn], fa[maxn];ll f[maxn], all[maxn];ll ans, t;bool flag[maxn], p[maxn];void dfs1(int x) {    f[x] = 0;    flag[x] = true;    ll temp = 0, son = 0;    if (fa[a[x]] != 0) s[fa[a[x]]].push_back(x);    else c[a[x]].push_back(x);    int k = fa[a[x]];    fa[a[x]] = x;    for (int i = 0; i < G[x].size(); i++) if (!flag[G[x][i]]) {        s[x].clear();        dfs1(G[x][i]);        f[x] += f[G[x][i]];        t = f[G[x][i]];        for (int j = 0; j < s[x].size(); j++) t -= f[s[x][j]];        all[a[x]] += 1LL*t*(t-1)/2;    }    f[x]++;    fa[a[x]] = k;}int main() {    //freopen("input.txt","r",stdin);    int cases = 0;    while (scanf("%d", &n) != EOF) {        cases++;        memset(p, 0, sizeof(p));        for (int i = 1; i <= n; i++) {            scanf("%d", &a[i]);            c[a[i]].clear();            all[a[i]] = 0;            p[a[i]] = true;        }        for (int i = 1; i <= n; i++) G[i].clear();        for (int i = 1; i < n; i++) {            scanf("%d %d", &x, &y);            G[x].push_back(y);            G[y].push_back(x);        }        memset(flag, 0, sizeof(flag));        ans = 0;        dfs1(1);        for (int i = 1; i <= n; i++) if (c[i].size() != 0) {            t = n;            for (int j = 0; j < c[i].size(); j++) t -= f[c[i][j]];            if (t > 0) all[i] += 1LL*t*(t-1)/2;        }        for (int i = 1; i <= n; i++) if (p[i]) {            ans += 1LL*n*(n-1)/2-all[i];        }        printf("Case #%d: %lld\n", cases, ans);    }}