[CF600E] Lomsat gelral

Description

给定一棵有 (n) 个结点的数，以 (1) 为根，树上的点的颜色为 (c[i])，对每个点，求以其为根的子树内的所有众数的和。

Solution

dsu on tree，首先预处理出每个点的重儿子 (wson[p])。然后 dfs 整棵树，处理 p 的所有孩子时，先跳过重孩子，最后处理重孩子时，保留其贡献，并暴力加上轻孩子们的贡献，最后得到总贡献。

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

struct Counter
{
    int c[N];

    vector<int> h;
    int ans = 0;
    int cnt = 0;
    void add(int p)
    {
        c[p]++;
        if (c[p] > ans)
            ans = c[p], cnt = p;
        else if (c[p] == ans)
            cnt += p;
        h.push_back(p);
    }
    int get()
    {
        return cnt;
    }
    void clear()
    {
        for (auto i : h)
            c[i] = 0;
        h.clear();
        ans = 0;
        cnt = 0;
    }
} counter;

int n;
vector<int> g[N];
int a[N], vis[N], wson[N], siz[N], ans[N];

void calc(int p, int fa)
{
    vis[p] = 1;
    siz[p] = 1;
    for (auto q : g[p])
    {
        if (q != fa)
        {
            calc(q, p);
            siz[p] += siz[q];
            if (siz[q] > siz[wson[p]])
                wson[p] = q;
        }
    }
}

void add(int p, int fa)
{
    counter.add(a[p]);
    for (auto q : g[p])
    {
        if (q != fa)
        {
            add(q, p);
        }
    }
}

void dfs(int p, int fa)
{
    vis[p] = 1;
    for (auto q : g[p])
    {
        if (q != fa && q != wson[p])
        {
            dfs(q, p);
            counter.clear();
        }
    }
    if(wson[p]) dfs(wson[p], p);
    for (auto q : g[p])
    {
        if (q != fa && q != wson[p])
        {
            add(q, p);
        }
    }
    counter.add(a[p]);
    ans[p] = counter.get();
}

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> a[i];

    for (int i = 1; i < n; i++)
    {
        int t1, t2;
        cin >> t1 >> t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }

    calc(1, 0);
    dfs(1, 0);

    for (int i = 1; i <= n; i++)
    {
        cout << ans[i] << " ";
    }
}