[HDU]-6035 Colorful Tree

URL : http://acm.hdu.edu.cn/showproblem.php?pid=6035

给出 $n$ 个点的树和各点颜色，$n<2\cdot 10^5$

求所有点与点 $\frac{n(n-1)}{2}$ 条路径上不同颜色数量的和

#include<stdio.h>
typedef long long LL;
const int MAXN = 2e5 + 5;
int n, color[MAXN], dot[MAXN<<1], lnk[MAXN<<1] ,nex[MAXN<<1] ,POS;
int up[MAXN], not_in[MAXN], not_in_root[MAXN];
LL ans;

inline LL getsum(int n){
    return 1LL * n * (n-1) >> 1;
}
inline void add_edge(int from, int to){
    dot[POS] = to;
    nex[POS] = lnk[from];
    lnk[from] = POS++;
}

LL dfs(int u, int fa) {
    int up_fa = up[color[u]];
    up[color[u]] = u;
    LL node_num = 1;
    for(int i = lnk[u]; i; i = nex[i]) {
        int &v = dot[i];
        if(v == fa) continue ;
        not_in[u] = 0;
        LL son_node_num = dfs(v,u);
        ans -= getsum(son_node_num - not_in[u]);
        node_num += son_node_num;
    }
    (up_fa ? not_in[up_fa] : not_in_root[color[u]]) += node_num;
    up[color[u]] = up_fa;
    return node_num;
}

int main()
{
    for(int Case = 1;scanf("%d", &n) != EOF; ++Case) {
        for(int i = 1; i <= n; ++i) {
            scanf("%d",color + i);
            not_in_root[i] = lnk[i] = 0;
        }
        POS = 1;
        for(int i = 1; i < n; ++i) {
            int x, y;
            scanf("%d%d", &x, &y);
            add_edge(x, y);
            add_edge(y, x);
        }
        ans = getsum(n) * n;
        dfs(1, -1);
        for(int i = 1;i <= n; ++i) {
            ans -= getsum(n - not_in_root[i]);
        }
        printf("Case #%d: %lld\n",Case,ans);
    }
    return 0;
}