[Codeforces600E] Lomsat gelral（树上启发式合并）

[Codeforces600E] Lomsat gelral（树上启发式合并）

题面

给出一棵N个点的树，求其所有子树内出现次数最多的颜色编号和。如果多种颜色出现次数相同，那么编号都要算进答案
N≤100000

分析

树上启发式合并，用map记录颜色出现次数，合并的时候更新最多的出现次数和编号和。

注意合并时的下标问题。当我们merge(x,y)的时候，由于是启发式合并，s[x]可能会并到s[y]去，如果我们直接查询s[x]，就查不到真正的答案。所以要再建立一个数组id[x],记录x的map合并之后到了哪里

代码

//树上启发式合并 
#include<iostream>
#include<cstdio>
#include<set>
#include<map>
#include<cstring>
#include<vector>
#define maxn 100000
using namespace std;
int n;
int c[maxn+5];
vector<int>E[maxn+5];

int id[maxn+5];
//当我们merge(x,y)的时候，由于是启发式合并，s[x]可能会并到s[y]去
//然后直接查询s[x]就炸了，所以id[x]就是记录x的map合并之后到了哪里 
map<int,int>s[maxn+5];
int ans_max[maxn+5];//最多的出现次数 
long long ans_cnt[maxn+5];//编号和 
long long res[maxn+5];
void merge(int x,int y){
    if(s[id[x]].size()<s[id[y]].size()){
        swap(id[x],id[y]);
    }
    int u=id[x],v=id[y];
    for(auto it : s[v]){
        int val=it.first;
        int cnt=it.second;
        s[u][val]+=cnt;
        if(s[u][val]>ans_max[u]){
            ans_max[u]=s[u][val];
            ans_cnt[u]=val; 
        }else if(s[u][val]==ans_max[u]){
            ans_cnt[u]+=val;
        }
    } 
}

void dfs(int x,int fa){
    for(int y : E[x]){
        if(y!=fa){
            dfs(y,x);
            merge(x,y);
        }
    }
    res[x]=ans_cnt[id[x]];
}

int main(){
    int u,v;
    scanf("%d",&n);
    for(int i=1;i<=n;i++){
        scanf("%d",&c[i]);
        id[i]=i;
        s[i][c[i]]=1;
        ans_max[i]=1;
        ans_cnt[i]=c[i];
    } 
    for(int i=1;i<n;i++){
        scanf("%d %d",&u,&v);
        E[u].push_back(v);
        E[v].push_back(u);
    }
    dfs(1,0);
    for(int i=1;i<=n;i++){
        printf("%I64d ",res[i]);
    }
}