[洛谷P2664] 樹上游戲
阿新 • • 發佈:2022-03-23
前言
本來想練一下點分樹來著,結果這題可以 \(O(n)\),結果還莫名拿了 rk1???
我看懂了題解的第一步轉化,然後死活看不懂接下來的步驟,被迫只能自己寫了/kk
題目
講解
拆貢獻是肯定的,但是這裡我們不拆點的貢獻,而是拆顏色的貢獻。
考慮對於一種顏色 col,我們把所有顏色是 col 的點相連的邊全部斷掉,可以發現每個點的貢獻是 n - 自己所在連通塊大小(當然顏色是 col 的點除外)。
令總顏色數量為 m,那麼每個點的答案就是:\(n\times m-\sum_{col}siz_{col,i}\),注意點的顏色是 col 時這個東西需要微調。
可以用樹上差分簡單求解,複雜度 \(O(n)\)
程式碼
//12252024832524 #include <bits/stdc++.h> #define TT template<typename T> using namespace std; typedef long long LL; const int MAXN = 100005; int n,m; int col[MAXN]; bool vis[MAXN]; LL ans[MAXN],dp[MAXN],wr[MAXN]; LL Read() { LL x = 0,f = 1; char c = getchar(); while(c > '9' || c < '0'){if(c == '-') f = -1;c = getchar();} while(c >= '0' && c <= '9'){x = (x*10) + (c^48);c = getchar();} return x * f; } TT void Put1(T x) { if(x > 9) Put1(x/10); putchar(x%10^48); } TT void Put(T x,char c = -1) { if(x < 0) putchar('-'),x = -x; Put1(x); if(c >= 0) putchar(c); } TT T Max(T x,T y){return x > y ? x : y;} TT T Min(T x,T y){return x < y ? x : y;} TT T Abs(T x){return x < 0 ? -x : x;} int head[MAXN],tot,hd[MAXN]; struct edge{ int v,nxt; }e[MAXN<<2]; void Add_Edge(int u,int v){ e[++tot] = edge{v,head[u]}; head[u] = tot; } void Add_Double_Edge(int u,int v){ Add_Edge(u,v); Add_Edge(v,u); } int f[MAXN],siz[MAXN],fff[MAXN]; void dfs1(int x,int fa) { siz[x] = 1; e[++tot] = edge{x,hd[f[col[x]]]}; hd[f[col[x]]] = tot; int tmp = f[col[x]]; f[col[x]] = x; for(int i = head[x],v; i ;i = e[i].nxt){ if((v = e[i].v) == fa) continue; LL lst = wr[x],val; dfs1(v,x); val = (!x ? 1ll * n * m : siz[v]) - (wr[x]-lst); dp[v] += val; //printf("add %d %lld\n",v,val); for(int j = hd[x]; j ;j = hd[x] = e[j].nxt){ if(!x) dp[e[j].v] -= n-fff[col[e[j].v]]; else dp[e[j].v] -= val; } siz[x] += siz[v]; } f[col[x]] = tmp; if(!f[col[x]]) fff[col[x]] += siz[x]; wr[f[col[x]]] += siz[x]; } void dfs2(int x,int fa){ ans[x] = dp[x] += dp[fa]; for(int i = head[x],v; i ;i = e[i].nxt){ if((v = e[i].v) == fa) continue; dfs2(v,x); } } int main() { // freopen(".in","r",stdin); // freopen(".out","w",stdout); n = Read(); for(int i = 1;i <= n;++ i){ col[i] = Read(); if(!vis[col[i]]) vis[col[i]] = 1,++m; } Add_Edge(0,1); for(int i = 1;i < n;++ i) Add_Double_Edge(Read(),Read()); dfs1(0,0); dfs2(1,0); for(int i = 1;i <= n;++ i) Put(1ll*n*m-ans[i]-n,'\n'); return 0; }