【CodeForces】600 E. Lomsat gelral (dsu on tree)

【题目】E. Lomsat gelral

【题意】给定n个点的树，1为根，每个点有一种颜色ci，一种颜色占领一棵子树当且仅当子树内没有颜色的出现次数超过它，求n个答案——每棵子树的占领颜色的编号和Σci（一棵子树可能有多种占领颜色）。1<=n,ci<=10^5。

【算法】dsu on tree

【题解】入门题，讲一下dsu on tree算法。

一、dsu on tree的适用范围：①子树询问，②支持数组上的快速信息加，③不带修。（注意不需要支持信息删除，只需要能清空信息）

如果写暴力的时候，每个点需要开一个O(n)的数组，然后要从下往上的合并数组，那么dsu on tree就可以帮你把复杂度降低到O(n log n)。

二、dsu on tree的过程：

①递归处理轻儿子，不保存信息。

②递归处理重儿子，保存信息。

③暴力加入该子树处理重儿子子树外的所有点的信息，得到该子树的答案。

④如果父边是轻边，扫描整棵子树清空信息。（操作必须和点数相关才能保证复杂度）

三、dsu on tree的复杂度证明：

算法思想是先处理轻儿子并且不保存，然后处理重儿子并且保存信息，这样每个点暴力扫描子树都不会包含重儿子子树。

轻重链剖分的特点是每个点到根至多经过log n条轻边，所以一般只要支持快速处理重链。（这个证明可以考虑最左端的底层节点，每跳一次轻边其右儿子都必须翻倍）

dsu on tree中每个点只会在其到根的路径中的所有轻边被遍历到1次，所以每个点的复杂度是O(log n)的。

总复杂度O(n log n)。

对于本题，需要支持快速信息加，记a[i]表示颜色i的出现次数，mx表示最多出现次数，sum表示答案，加的时候更新一下，清空的时候直接mx=sum=0即可。

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int maxn=100010;
int c[maxn],first[maxn],tot,sz[maxn],w[maxn],a[maxn],n,mx;
long long ans[maxn],sum;
bool vis[maxn];
struct edge{int v,from;}e[maxn*2];
void insert(int u,int v){tot++;e[tot].v=v;e[tot].from=first[u];first[u]=tot;}
void predfs(int x,int fa){
    sz[x]=1;
    for(int i=first[x];i;i=e[i].from)if(e[i].v!=fa){
        predfs(e[i].v,x);
        sz[x]+=sz[e[i].v];
        if(sz[e[i].v]>sz[w[x]])w[x]=e[i].v; 
    }
}
void add(int x,int fa,int k){
    a[c[x]]+=k;
    if(k>0){
        if(a[c[x]]>mx){
            mx=a[c[x]];sum=c[x];
        }
        else if(a[c[x]]==mx)sum+=c[x];
    }
    for(int i=first[x];i;i=e[i].from)if(e[i].v!=fa&&!vis[e[i].v])add(e[i].v,x,k);
}
void dfs(int x,int fa,int k){
    for(int i=first[x];i;i=e[i].from)if(e[i].v!=fa&&e[i].v!=w[x])dfs(e[i].v,x,0);
    if(w[x])dfs(w[x],x,1),vis[w[x]]=1;
    add(x,fa,1);
    if(w[x])vis[w[x]]=0;
    ans[x]=sum;
    if(!k)add(x,fa,-1),mx=sum=0;
}
int main(){
    scanf("%d",&n);
    for(int i=1;i<=n;i++)scanf("%d",&c[i]);
    for(int i=1;i<n;i++){
        int u,v;scanf("%d%d",&u,&v);
        insert(u,v);insert(v,u);
    }
    predfs(1,0);dfs(1,0,1);
    for(int i=1;i<=n;i++)printf("%lld ",ans[i]);
    return 0;
}