Educational Codeforces Round 2 E. Lomsat gelral（树上启发式合并）

题目链接
题意：

给出一颗有$n$个点的树，根节点为$1$，每个点都有一种颜色，子树 $i$ 的主导颜色为数量最多的颜色，假如有多个颜色数量最多，那么这些颜色都是主导颜色，所有子树的主导颜色之和。

题解：

典型的树上查询问题，很明显，不能暴力查找，但是可以"优雅"的暴力查找，即树上启发式合并。

先预处理出每个节点的重儿子，在查询子树 $i$ 的时候，先查询轻儿子，把贡献删除，再查询重子树，保留贡献，再暴力查找轻儿子算 $i$ 的贡献，因为在算重儿子的时候，贡献保留，所以此时并不用查找重儿子。最后再判断子树 $i$ 是否为重儿子，即是否要删除贡献。

代码：

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<queue>
#include<map>
#include<stack>
#include<set>
#include<ctime>
#define iss ios::sync_with_stdio(false)
using namespace std;
typedef unsigned long long ull;
typedef long long ll;
typedef pair<int,int>pii;
const int MAXN=1e5+5;
const int mod=1e9+7;
const double eps = 1e-6;
struct node
{
    int to;
    int next;
}e[MAXN<<1];
int head[MAXN];
int col[MAXN];
int son[MAXN];
int sz[MAXN],num[MAXN];
ll dp[MAXN];
int cnt;
int maxx,now;
ll ans;
void add(int u,int  v)
{
    e[cnt].to=v;
    e[cnt].next=head[u];
    head[u]=cnt++;
}
void dfs(int u,int f)//找重儿子
{
    sz[u]=1;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(v==f) continue;
        dfs(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]])
        {
            son[u]=v;
        }
    }
}
void cal(int u,int f,int val)
{
    num[col[u]]+=val;
    if(num[col[u]]==maxx)
    {
        ans+=col[u];
    }
    else if(num[col[u]]>maxx)
    {
        maxx=num[col[u]];
        ans=col[u];
    }
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(v==f||v==now) continue;
        cal(v,u,val);
    }
}
void solve(int u,int f,bool flag)
{
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(v==f||v==son[u]) continue;
        solve(v,u,false);
    }
    if(son[u]) solve(son[u],u,true),now=son[u];
    cal(u,f,1);
    now=0;
    dp[u]=ans;
    if(!flag)
    {
        cal(u,f,-1);
        ans=maxx=0;
    }
}
int main()
{
    memset(head,-1,sizeof head);
    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;
        add(u,v);
        add(v,u);
    }
    dfs(1,-1);
    solve(1,-1,true);
    for(int i=1;i<=n;i++)
    {
        printf("%lld ",dp[i]);
    }
}