Educational Codeforces Round 2 E. Lomsat gelral(dsu)

题目链接

题意:给你一棵以1为根n个点的树，问你以i为根的子树的众数和是多少

思路:dsu是一种优化暴力的手段 首先进行轻重链剖分 然后只记录重链的信息 轻链的信息就直接暴力查找 经过证明这样复杂度可以是nlogn。

#include<bits/stdc++.h>
#define ll long long int
using namespace std;
inline ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
inline ll lcm(ll a,ll b){return a/gcd(a,b)*b;}
int moth[13]={0,31,28,31,30,31,30,31,31,30,31,30,31};
int dir[4][2]={1,0 ,0,1 ,-1,0 ,0,-1};
int dirs[8][2]={1,0 ,0,1 ,-1,0 ,0,-1, -1,-1 ,-1,1 ,1,-1 ,1,1};
const int inf=0x3f3f3f3f;
const ll mod=1e9+7;
const int N=100007;
vector<int> G[N];
int col[N],dp[N],son[N];
ll cnt[N],ans[N];
ll maxx=-inf,sum=0;
int po;
void dfs(int u,int fa){
    dp[u]=1;
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i];
        if(v==fa) continue;
        dfs(v,u);
        dp[u]+=dp[v];
        if(dp[v]>dp[son[u]]) son[u]=v; //找重儿子
    }
}
void add(int u,int fa,int w){
    cnt[col[u]]+=w;
    if(cnt[col[u]]>maxx) maxx=cnt[col[u]],sum=col[u];
    else if(cnt[col[u]]==maxx) sum+=col[u];
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i];
        if(v==fa||v==po) continue;
        add(v,u,w);
    }
}
void dfss(int u,int fa,int opt){
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i];
        if(v==fa) continue;
        if(son[u]!=v) dfss(v,u,0); //优先遍历轻儿子
    }
    if(son[u]) dfss(son[u],u,1),po=son[u];//po是标记当重链的重儿子 在算当前节点时可以不用遍历重儿子
    add(u,fa,1);
    po=0;
    ans[u]=sum;
    if(!opt) add(u,fa,-1),maxx=-inf,sum=0; //消除轻儿子的影响
}
int main(){
    ios::sync_with_stdio(false);
    int n; cin>>n;
    for(int i=1;i<=n;i++){
        cin>>col[i];
    }
    for(int i=1;i<=n-1;i++){
        int u,v; cin>>u>>v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1,0); dfss(1,0,1);
    for(int i=1;i<=n;i++)
        cout<<ans[i]<<" ";
    return 0;
}