CF-600-E-启发式合并
600-E 题目大意
给定一颗\(n\)个节点的树,根为\(1\)。树上的每个节点\(i\)都有一个颜色\(c_i\)。
如果一个颜色在以\(x\)为根的子树中出现次数最多,那么称该颜色为主要颜色,显然,一颗树中可以有多个主要颜色。
求出对于每个节点为根时,其子树中所有主要颜色的编号和。
Solution
启发式合并思路:
在\(dfs\)的过程中,每个节点维护其子树中所有的颜色数量信息,主要颜色的出现的次数,主要颜色的编号和。
每次递归回溯时,用重儿子的节点信息来更新当前节点的信息,这里利用启发式合并更新信息即可,无需预处理出所有节点的重儿子,细节见代码。
时间复杂度\(O(nlogn)\)。
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int main(){
ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
int n;
cin>>n;
vector<vector<int>> e(n);
vector<int> mx(n);
vector<ll> sum(n);
map<int,int> p[n];
for(int i=0;i<n;i++){
int col;
cin>>col;
p[i][col]++;
sum[i]=col;
}
for(int i=1;i<n;i++){
int x,y;
cin>>x>>y;
--x,--y;
e[x].push_back(y);
e[y].push_back(x);
}
function<void(int,int)> dfs=[&](int x,int fa){
mx[x]=1;
for(auto y:e[x]){
if(y==fa) continue;
dfs(y,x);
if(p[x].size()<p[y].size()){
swap(p[x],p[y]);
mx[x]=mx[y];
sum[x]=sum[y];
}
for(auto &[a,b]:p[y]){
p[x][a]+=b;
if(p[x][a]>mx[x]){
mx[x]=p[x][a];
sum[x]=a;
}else if(p[x][a]==mx[x]){
sum[x]+=1LL*a;
}
}
}
};
dfs(0,-1);
for(int i=0;i<n;i++){
cout<<sum[i]<<" ";
}
return 0;
}