hdu 6035 Colorful Tree

Colorful Tree

题目

题意
给定一颗树，再给定树上每一个节点一个颜色，求树上总共n(n-1)/2条路径每条路径上不同颜色数的总和。

做法
首先，

$$总和=总路径数 * 颜色数 - \sum_{每一种颜色}^{}没参与的贡献$$ .

对于某一个id节点来说，算没参与贡献，可以通过求其最高同色子节点，以这些节点和其子树节点与id形成的路径必然会有重复id节点的颜色，即id节点对这些路径没有贡献。
也就是求对于id节点这个颜色来说，它的同色子节点对其子树节点造成的最大“截断”数量。
最后还要减去所有颜色剩下的连通块。

#include <bits/stdc++.h>

using namespace std;
const int maxn = 200005;
typedef long long ll; 
int n;
int c[maxn],T,tot;
int size[maxn];
bool vis[maxn];
ll ans;
vector<int> e[maxn];
ll s[maxn];

void dfs(int id,int fa){
    size[id]=1;
    ll son=0;
    for(auto ep:e[id]){
        if(ep==fa) continue;
        ll last = s[c[id]];
        dfs(ep,id);
        size[id] += size[ep];
        ll add = s[c[id]] - last;
        ans -= (ll)(size[ep] - add) * (size[ep] - add - 1LL) / 2LL;
        son += size[ep] - add;
    }
    s[c[id]] += son + 1;
}

int main(){
    freopen("input.txt","r",stdin);
    int T=1;
    while(~scanf("%d",&n)){
        for(int i=0;i<=n;i++)
            e[i].clear();
        memset(vis,0,sizeof vis);
        memset(s,0,sizeof s);
        ll cnt = 0;
        for(int i=1;i<=n;i++){
            scanf("%d",&c[i]);
            if(!vis[c[i]]) cnt++;
            vis[c[i]]=1;
        }
        for(int i=1;i<n;i++){
            int u,v;
            scanf("%d%d",&u,&v);
            e[u].push_back(v);
            e[v].push_back(u);
        }
        ans = (ll)n * (n-1LL) / 2LL * cnt;
        dfs(1,0);
        for(int i=1;i<=n;i++){
            if(!vis[i]) continue;
            ans -= (ll)(n - s[i]) * (n - s[i] - 1LL) / 2LL;
        }
        printf("Case #%d: %lld\n",T++,ans);
    }
    return 0;
}