[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 }