CF600E Lomsat gelral

Lomsat gelral - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

Solution

树上启发式合并模板题。

考虑暴力做法：设以 $i$ 为根的子树的答案为 $A_i$，每次求 $A_i$ 都重新遍历整棵子树。选择从 $1$ 开始往下 dfs，然后从子结点开始一个一个地求 $A_i$，时间复杂度为 $O(n^2)$。

定义与 $i$ 父亲相同（同为 $f_i$）且深度相同的结点为 $i$ 的兄弟结点，设为 $j$。不妨设 $A_i$ 的求解在 $A_j$ 前面，则求 $A_j$ 时，需要把求 $A_i$ 时所用的计数器等变量重新清零。

注意到作为 $f_i$ 的儿子，$A_i$ 和 $A_j$ 都与 $A_{f_i}$ 有关联。显然最后一个求 $A_i$ 的 $f_i$ 的子结点在求完 $A_i$ 后计数器不用清零，可直接用于 $A_{f_i}$ 的计算中。那么我们选择 $f_i$ 的重儿子作为最后一个求 $A_i$ 的点。

换句话说，计算 $A_i$ 时，不用遍历 $i$ 的重儿子的子树，轻儿子的子树要全部遍历。优化后，时间复杂度变为 $O(n\log n)$。具体证明可见这里。

Code

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

const ll N = 1e5 + 10;

ll n, c[N], siz[N], son[N], ans[N], cnt[N], maxn, sum;
vector<ll> g[N];

void clear(ll u, ll fa)
{
    cnt[c[u]] = 0;
    for (ll v : g[u]) {
        if (v == fa) continue;
        clear(v, u);
    }
}

void calc(ll u, ll fa, ll s)
{
    cnt[c[u]]++;
    if (cnt[c[u]] > maxn) {
        maxn = cnt[c[u]];
        sum = c[u];
    } else if (cnt[c[u]] == maxn)
        sum += c[u];
    for (ll v : g[u]) {
        if (v == fa || v == s) continue;
        calc(v, u, s);
    }
}

void dfs1(ll u, ll fa)
{
    siz[u] = 1;
    for (ll v : g[u]) {
        if (v == fa) continue;
        dfs1(v, u);
        if (siz[v] > siz[son[u]])
            son[u] = v;
        siz[u] += siz[v];
    }
}

void dfs2(ll u, ll fa)
{
    for (ll v : g[u]) {
        if (v == fa || v == son[u]) continue;
        dfs2(v, u);
        maxn = sum = 0;
        clear(v, u);
    }
    if (son[u]) 
        dfs2(son[u], u);
    calc(u, fa, son[u]);
    ans[u] = sum;
}

int main()
{
    scanf("%lld", &n);
    for (int i = 1; i <= n; i++)
        scanf("%lld", &c[i]);
    for (int i = 1; i < n; i++) {
        ll x, y;
        scanf("%lld%lld", &x, &y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    dfs1(1, 0); dfs2(1, 0);
    for (int i = 1; i <= n; i++)
        printf("%lld ", ans[i]);
    return 0;
}