HDU6035-Colorful Tree-虛樹思想

yoshinow2001發表於2024-08-19

link：https://codeforces.com/gym/102253/problem/C
題意：一棵樹，每個點有顏色，求所有路徑上出現的顏色個數之和，基本要求線性。

對每種顏色考慮答案，列舉到顏色c的時候想著把所有顏色c的點拿出來，建立虛樹，這樣總的點數是 \(O(n)\) 的，但建虛樹其實要 \(O(n\log n)\) 的時間，而且建完了的DP似乎也很麻煩。

對這個虛樹考慮，如果顏色c用黑點表示，那麼假設刪去黑點，樹會變成若干連通塊，這個顏色的貢獻就是 \(\binom{n}{2}-\sum_{block}\binom{sz}{2}\)

考慮從上往下模擬這個dfs的過程，對當前的結點x，對每個孩子結點 \(to\)，求出 \(sz[to]-\sum\) 從上往下\(c[x]\) 第一次出現的那些點的子樹大小，就能確定對應的連通塊大小，然後統計 \(x\) 的貢獻。
這個子樹大小可以透過一些類似樹上差分的技巧完成：dfs時記錄每種顏色當前還“剩下”幾個點，一開始全部設為 \(n\) 個點，每次搜完某個 \(x\) ，就將其子樹內的點全部刪去——當然 \(x\) 內可能還會有 \(c[x]\) 顏色的其他點，稍微算一下差值就好

#include<bits/stdc++.h>
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define endl '\n'
#define fastio ios_base::sync_with_stdio(false);cin.tie(0);cout.tie(0)
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int N=2e5+5;
int n,c[N],sz[N],cnt[N];
vector<vector<int>> G;
ll ans;
ll C2(int n){return (ll)n*(n-1)/2;}
void dfs1(int x,int fa){
    sz[x]=1;
    for(auto to:G[x])if(to!=fa){
        dfs1(to,x);
        sz[x]+=sz[to];
    }
}
void dfs2(int x,int fa){
    if(x==1)rep(i,1,n)cnt[i]=n;
    int cur=cnt[c[x]];
    for(auto to:G[x])if(to!=fa){
        int before=cnt[c[x]];
        dfs2(to,x);
        int after=cnt[c[x]];
        int del=sz[to]-(before-after);
        ans+=C2(del);
        cnt[c[x]]-=del;
    }
    int lst=cnt[c[x]];
    cnt[c[x]]-=sz[x]-(cur-lst);
    if(x==1)rep(i,1,n)ans+=C2(cnt[i]);
}
int main(){
    fastio;
    int tc=0;
    while(cin>>n){
        tc++;
        rep(i,1,n)cin>>c[i];
        G=vector<vector<int>>(n+1);
        ans=0;
        rep(i,1,n-1){
            int u,v;
            cin>>u>>v;
            G[u].push_back(v);
            G[v].push_back(u);
        }
        dfs1(1,-1);
        dfs2(1,-1);
        cout<<"Case #"<<tc<<": "<<n*C2(n)-ans<<endl;
    }
    return 0;
}

相關文章