hdu6035

201707 （ACM）想法

• 每个颜色的贡献 = 所有路径数 - 不经过这个颜色的贡献
• 切掉某颜色，就有很多联通块。
• 如何求联通块的大小？
  
  • dfs后序。
  • 维护一些信息。假设当前到 $u$ 点。
  • 子树的信息 = 遍历子树回来 - 刚到 $u$
  • 假设当前点 $v$ ，父亲是 $u$，那么 [ $u$切割成的，包含$v$ ] 的联通块大小为 $sz[v]-sz[v子树中颜色为col[u]的]$
  • 所以维护的信息 $sum[i]$ 表示已经存的信息中以 $i$ 颜色为根的子树大小（如果当前子树 $i$ 的子孙与 $i$ 同色，不会重复计算）
• 总结：这题中有关dfs的部分还是很值得学习的。

#include<bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define pb push_back
#define sz(a) (int)a.size()
#define mp make_pair
#define rep(i, a, b) for(int i=(a); i<(b); i++)
#define de(x) cout<<#x<<"="<<x<<endl
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;
//--------
const int N=200005;
bool vis[N];
int n;
int col[N],tsz[N],sum[N];
ll ans;
vector<int> G[N];
void dfs(int u,int f) {
    tsz[u]=1;
    int tt=0;
    rep(i,0,sz(G[u])) {
        int v=G[u][i];
        if(v==f) continue;
        ///
        int res=-sum[col[u]];
        dfs(v,u);
        res+=sum[col[u]];
        tt+=res;
        res=tsz[v]-res;
        ans=ans-1ll*res*(res-1)/2;
        ///
        tsz[u]+=tsz[v];
    }
    sum[col[u]]+=tsz[u]-tt;
}
int main() {
    int ca=0;
    while(~scanf("%d",&n)) {
        ///init
        rep(i,0,n+1) G[i].clear();
        memset(vis,0,sizeof(vis));
        memset(sum,0,sizeof(sum));
        memset(col,-1,sizeof(col));
        ///read
        int cntc=0;
        rep(i,1,n+1) {
            scanf("%d",col+i);
            if(!vis[col[i]]) {
                vis[col[i]]=1;
                ++cntc;
            }
        }
        rep(i,1,n) {
            int u,v;scanf("%d%d",&u,&v);
            G[u].pb(v);
            G[v].pb(u);
        }
        ///solve
        ans=1ll*n*(n-1)/2*cntc;
        dfs(1,0);
        rep(i,1,n+1) {
            if(vis[i]) {
                int res=n-sum[i];
                ans=ans-1ll*res*(res-1)/2;
            }
        }
        printf("Case #%d: %lld\n",++ca,ans);
    }
    return 0;
}