树上启发式合并
题意:
一棵树每个节点有一种颜色,某个节点的主导颜色为该子树下颜色数量最多的颜色编号,求每个节点的主导颜色和 思路:
- 每遍历完一颗子树的点,就要把最大值清空(避免影响另一颗子树的计算)
- 当算完所有子节点后,就要统计当前点,需要在遍历一次子节点将颜色加回来
- 可以发现遍历的最后一棵子树不用清空
- 将重儿子作为最后一次遍历的子树
- 每个轻儿子被遍历的次数为到根节点轻边的数
- 对于一条轻边,$ 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];
}