E. Lomsat gelral (dsu on tree)

题目

You are given a rooted tree with root in vertex 1. Each vertex is coloured in some colour. Let's call colour c dominating in the subtree of vertex v if there are no other colours that appear in the subtree of vertex v more times than colour c. So it's possible that two or more colours will be dominating in the subtree of some vertex. The subtree of vertex v is the vertex v and all other vertices that contains vertex v in each path to the root. For each vertex v find the sum of all dominating colours in the subtree of vertex v.

输入

The first line contains integer n (1 ≤ n ≤ 105) — the number of vertices in the tree. The second line contains n integers ci (1 ≤ ci ≤ n), ci — the colour of the i-th vertex. Each of the next n - 1 lines contains two integers xj, yj (1 ≤ xj, yj ≤ n) — the edge of the tree. The first vertex is the root of the tree.

解法

这是一道树上启发式合并的板子题我们先用第一遍dfs求出每个结点的重儿子，然后第二遍dfs的时候对于轻儿子我们消除他的贡献对于重儿子我们暴力保留他的贡献然后输出答案即可

Code

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define int long long
#define endl '\n'
#define pb push_back
#define NO cout << "NO" << endl;
#define YES cout << "YES" << endl;
#define fi first
#define se second
#define all(x) (x).begin(),(x).end()
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define per(i,a,n) for (int i=n;i>=a;i--)
typedef vector<int> VI;
typedef pair<int,int> PII;
typedef vector<VI> VII;
ll MOD;
ll powmod(ll a,ll b) {ll res=1;a%=MOD; assert(b>=0); for(;b;b>>=1){if(b&1)res=res*a%MOD;a=a*a%MOD;}return res;}
mt19937 mrand(random_device{}()); 
const int N = 2e5 + 10;
void solve()
{
    int n; cin >> n;
    VI a(n + 1); rep (i, 1, n) cin >> a[i];
    VII e(n + 1);
    rep (i, 2, n) {
        int u, v; cin >> u >> v;
        e[u].pb(v);
        e[v].pb(u);
    }
    VI son(n + 1), sz(n + 1), vis(n + 1), res(n + 1), cnt(n + 1);
    function<void(int, int)> dfs1 = [&] (int u, int fa) {
        sz[u] = 1;
        for (auto v : e[u]) {
            if (v == fa) continue;
            dfs1(v, u);
            sz[u] += sz[v];
            if (!sz[son[u]] ||sz[v] > sz[son[u]]) {
                son[u] = v;
            }
        }
    };
    int mx = 0, flag = 0, sum = 0;
    function<void(int, int, int)> count = [&] (int u, int fa, int val) {
        cnt[a[u]] += val;
        if (cnt[a[u]] > mx) {
            mx = cnt[a[u]];
            sum = a[u];
        } else if (cnt[a[u]] == mx) {
            sum += a[u];
        }
        for (auto v : e[u]) {
            if (v == fa || v == flag) continue;
            count(v, u, val);
        }
    };
    function<void(int, int, bool)> dfs2 = [&] (int u, int fa, bool keep) {
        for (auto v : e[u]) {
            if (v == fa || v == son[u]) continue;
            dfs2(v, u, false);
        }
        if (son[u]) {
            dfs2(son[u], u, true);
            flag = son[u];
        }
        count(u, fa, 1);
        flag = 0;
        res[u] = sum;
        if (!keep) {
            count(u, fa, -1);
            sum = mx = 0;
        }
    };
    dfs1(1, 0);
    dfs2(1, 0, false);
    rep (i, 1, n) cout << res[i] << " \n"[ i == n];

}
signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    // int T;cin >> T;
    // while ( T -- )
    solve();
    return 0;
}