hdu6035 Colorful Tree

题解：这个图构成了一树，单独考虑每一种颜色，答案就是对于每种颜色至少经过一次这种的路径条数之和。反过来思考只需要求有多少条路径没有经过这种颜色即可。然后没有经过路线的条数就是以这个颜色为边界构成联通块能够组成的路径数目，路径数目就是联通块个数*（联通块个数-1）/2

而这个点到下面结点的联通块个数就是这个颜色子节点为根的个数减去下面跟这个颜色一样的节点并且以这个节点为根树的节点个数总和
如何求下面代码有注释

#include<iostream>
#include<cstring>
#include<algorithm>
#include<cstdio>
#include<vector>
using namespace std;
const int mx = 2e5+5;
typedef long long int ll;
vector<int>g[mx];
int a[mx];
ll siz[mx];
ll sum[mx];
bool col[mx];
ll de;
void dfs(int u,int fa){
    siz[u] = 1;
    ll pre = sum[a[u]]; //记录上次还没有往下搜的时候已经这个颜色的点被覆盖了多少
    ll ans = 0;
    for(auto v: g[u]){
        if(v==fa)   continue;
        dfs(v,u);
        siz[u]+=siz[v];
        ll cnt = siz[v]-sum[a[u]]+pre;   //下面搜过去下面的节点数减去增加的颜色覆盖就是联通块的大小，当然如果下面没有这个sum[a[u]]是不会增加的
        de += cnt*(cnt-1)/2;             //cnt就是以这个点为边界求出来的连通快大小
        ans += cnt;                      //答案要加上下面这个点覆
        pre = sum[a[u]];
    }
    sum[a[u]]+=ans+1;   //因为本身这个点也算一个点所有要加一
}
int main(){
    int n,u,v;
    int casei = 1;
    while(scanf("%d",&n)!=EOF){
        memset(col,0,sizeof(col));
        memset(siz,0,sizeof(siz));
        memset(sum,0,sizeof(sum));
        for(int i = 1; i <= n; i++){
            scanf("%d",&a[i]);
            col[a[i]] = 1;
            g[i].clear();
        }
        for(int i = 1; i < n; i++){
            scanf("%d%d",&u,&v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        ll num = 0;
        de = 0;
        for(int i = 1; i <= n; i++) num += col[i];
        ll ans =  num*n*(n-1)/2;  //假设每种颜色每条道路上都有
        dfs(1,0);
        ans -= de;
        for(int i = 1; i <= n; i++){
            if(i!=a[1]&&col[i]){
                ll cnt = n-sum[i];    //还未贡献的值就是剩下几个最高结点之间的连通快
                ans -= cnt*(cnt-1)/2;    
            }
        }
        printf("Case #%d: %I64d\n",casei++,ans);
    }
    return 0;
}