传送门:https://codeforces.com/problemset/problem/600/E
题意:
给一颗n个点的有根树,每个点有一种颜色,一颗子树中出现次数最多的颜色称为主导颜色,主导颜色可能有多个,求这棵树的每个结点的对应子树的主导颜色之和
思路:
DSU on tree 树上启发式合并,先记录一下板子
··········································································································
先树链剖分找重儿子,重儿子就是最大子树的根节点
dfs每个节点,暴力计算所有轻儿子的贡献,并合并到重儿子上去
注意,轻儿子的贡献在计算之后必须要删除,不然可能会MLE,
如果都不保留,可能会TLE,所以需要平衡一下时间和空间,我们保留重儿子的贡献即可,这也是DSU on tree的意义
时间复杂度O(nlogn),以下分析来自Yveh的[trick]dsu on tree
在这种做法中,我们先进行树链剖分。
dfs的时候,首先dfs节点x的轻儿子,暴力消去影响,再dfs节点x的下一个轻儿子,依此类推。
然后dfs节点x的重儿子,无需消去影响。
在最后,我们为了统计x,再将x轻儿子的贡献加回来。
看起来很暴力,但是实际上它的时间复杂度是O(nlogn)的,跑得飞快。
可以这么考虑:
只有dfs到轻边时,才会将轻边的子树中合并到上一级的重链,树链剖分将一棵树分割成了不超过logn条重链。
每一个节点最多向上合并logn次,单次修改复杂度O(1)。
所以整体复杂度是O(nlogn)的。
AC代码:
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+10;
struct Edge {
int to,next;
} e[maxn<<1];
int n,u,v,tot=1,head[maxn],siz[maxn],son[maxn];
int col[maxn],cnt[maxn],flag,maxc;
ll ans[maxn],sum;
inline void addEdge(int u,int v) {
e[tot].to = v;
e[tot].next = head[u];
head[u]=tot++;
}
//重链剖分找重儿子
void dfs(int u,int f) {
siz[u]=1;
for(int i=head[u]; i; i=e[i].next) {
int v=e[i].to;
if(v==f) continue;
dfs(v,u);
siz[u]+=siz[v];
if(siz[v]>siz[son[u]]) son[u]=v;
}
}
//统计某结点及其所有轻儿子的贡献
void count(int u,int f,int val) {
cnt[col[u]] += val;//val为正负用来控制是加贡献还是删除贡献
if(cnt[col[u]]>maxc) {
maxc=cnt[col[u]];
sum=col[u];
} else if(cnt[col[u]]==maxc) { //两个颜色数量相同都是主导颜色都得算
sum+=col[u];
}
//排除被标记的重儿子,统计其他儿子子树信息
for(int i=head[u]; i; i=e[i].next) {
int v=e[i].to;
if(v==f || v==flag) continue;
count(v,u,val);
}
}
//DSU on tree 的板子
void dfs(int u,int f,bool keep) {
//第一步:搞轻儿子及其子树 算其答案 删贡献
for(int i=head[u]; i; i=e[i].next) {
int v=e[i].to;
if(v==f || v==son[u]) continue;
dfs(v,u,false);
}
//第二步:搞重儿子及其子树算其答案 不删贡献
if(son[u]) {
dfs(son[u],u,true);
flag=son[u];
}
//第三步:暴力统计u及其所有轻儿子的贡献合并到刚算出的重儿子信息里
count(u,f,1);
flag = 0;
ans[u] = sum;
//把需要删除贡献的删一删
if(!keep) {
count(u,f,-1);
sum=maxc=0;
}
}
int main() {
cin>>n;
for(int i=1; i<=n; i++) cin>>col[i];
for(int i=1; i<n; i++) {
cin>>u>>v;
addEdge(u,v),addEdge(v,u);
}
dfs(1,0);
dfs(1,0,0);
for(int i=1; i<=n; i++) {
cout<<ans[i]<<" ";
}
}