Skip to content

树上启发式合并 ​

题意:
一棵树每个节点有一种颜色,某个节点的主导颜色为该子树下颜色数量最多的颜色编号,求每个节点的主导颜色和 思路:

  • 每遍历完一颗子树的点,就要把最大值清空(避免影响另一颗子树的计算)
  • 当算完所有子节点后,就要统计当前点,需要在遍历一次子节点将颜色加回来
  • 可以发现遍历的最后一棵子树不用清空
  • 将重儿子作为最后一次遍历的子树
  • 每个轻儿子被遍历的次数为到根节点轻边的数
  • 对于一条轻边,$ siz_{重儿子}>=siz_{轻儿子} $ ,因此轻边最多$ logn $ 条
  • 时间复杂度: $ O(nlogn) $
cpp
int color[N],siz[N],son[N],cnt[N];
vector<int> adj[N];
ll ans[N];
ll sum,mx,n;
void dfs_son(int u,int fa){
    siz[u]=1;
    for(auto v:adj[u]){
        if(v==fa)continue;
        dfs_son(v,u);
        siz[u]+=siz[v];
        if(siz[v]>siz[son[u]])son[u]=v;
    }
}
void update(int u,int fa,int sign,int pson){//将u子树的颜色加(减)到集合
    int c=color[u];
    cnt[c]+=sign;
    if(cnt[c]>mx)mx=cnt[c],sum=c;
    else if(cnt[c]==mx)sum+=c;
    for(auto v:adj[u]){
        if(v==fa||v==pson)continue;
        update(v,u,sign,pson);
    }
}
void dfs(int u,int fa,int op){
    for(auto v:adj[u]){
        if(v==fa||v==son[u])continue;
        dfs(v,u,0);
    }
    if(son[u])dfs(son[u],u,1);
    update(u,fa,1,son[u]);//更新除重儿子树外的点
    ans[u]=sum;
    if(!op)update(u,fa,-1,0),sum=mx=0;//当前节点为轻儿子,所有节点删除
}
void solve(){
    cin>>n;
    for(int i=1;i<=n;i++)cin>>color[i];
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs_son(1,0);
    dfs(1,0,1);
    for(int i=1;i<=n;i++)cout<<ans[i]<<" \n"[i==n];
}