[CF600E]Dsu on tree
题意:树上每个点都有颜色,称一个颜色占领一棵子树,当且仅当没有别的颜色在这棵子树内的数量比它多。求所有子树的占领颜色之和。题解:最显然的是DFS序+主席树或莫队,这里使用Dsu on tree。
每次暴力DFS之后,只撤销除重儿子之外的点的贡献。由于重儿子的性质,均摊后复杂度为$O(n\log n)$。
1 #include<cstdio> 2 #include<algorithm> 3 #include<iostream> 4 #define rep(i,l,r) for (int i=(l); i<=(r); i++) 5 #define For(i,x) for (int i=h[x],k; i; i=nxt[i]) 6 typedef long long ll; 7 using namespace std; 8 9 const int N=1000010; 10 int n,u,v,cnt,tot[N],mx,col[N],sz[N],son[N],h[N],to[N],nxt[N]; 11 ll ans[N],sm; 12 bool skip[N]; 13 14 void add(int u,int v){ to[++cnt]=v; nxt[cnt]=h[u]; h[u]=cnt; } 15 16 void get(int x,int fa){ 17 sz[x]=1; 18 For(i,x) if ((k=to[i])!=fa){ 19 get(k,x); sz[x]+=sz[k]; 20 if (sz[k]>sz[son[x]]) son[x]=k; 21 } 22 } 23 24 void dfs(int x,int fa,int op){ 25 tot[col[x]]+=op; 26 if (op>0 && tot[col[x]]>=mx){ 27 if (tot[col[x]]>mx) sm=0,mx=tot[col[x]]; 28 sm+=col[x]; 29 } 30 For(i,x) if ((k=to[i])!=fa && !skip[k]) dfs(k,x,op); 31 } 32 33 void work(int x,int fa,bool cl){ 34 For(i,x) if ((k=to[i])!=fa && k!=son[x]) work(k,x,1); 35 if (son[x]) work(son[x],x,0),skip[son[x]]=1; 36 dfs(x,fa,1); ans[x]=sm; skip[son[x]]=0; 37 if (cl) dfs(x,fa,-1),mx=sm=0; 38 } 39 40 int main(){ 41 scanf("%d",&n); 42 rep(i,1,n) scanf("%d",&col[i]); 43 rep(i,2,n) scanf("%d%d",&u,&v),add(u,v),add(v,u); 44 get(1,0); work(1,0,0); 45 rep(i,1,n) cout<<ans[i]<<' '; 46 return 0; 47 }