题解：洛谷 CF600E Lomsat gelral

【题目来源】

洛谷：CF600E Lomsat gelral

【题目描述】

给你一棵以结点 \(1\) 为根的有根树，每个节点最开始都被涂上了颜色。

如果颜色 \(c\) 在以结点 \(v\) 为根的子树中出现次数最多，则称其在以结点 \(v\) 为根的子树中占重要地位。一棵树中可以有很多颜色同时占重要地位。

以 \(v\) 为根的子树指结点 \(v\) 及其他到根结点的路径包含 \(v\) 的结点。

请输出对于每一个结点 \(v\)，在其子树中占重要地位的颜色编号之和。

【输入】

第一行包含一个整数 \(n\)，表示树的结点个数。

第二行包含 \(n\) 个整数 \(c _ i\)，表示每个结点的颜色。

接下来 \(n - 1\) 行，每行包含两个整数 \(x _ j\) 和 \(y _ j\)，表示每条边。结点 \(1\) 是树的根。

【输出】

输出一行 \(n\) 个整数，表示对于每个节点，在以其为根的子树中占重要地位的颜色编号之和。

【输入样例】

4
1 2 3 4
1 2
2 3
2 4

【输出样例】

10 9 3 4

【核心思想】

1. 问题分析：给定一棵以 \(1\) 为根的有根树，每个节点有一种颜色。对每个节点 \(v\)，需要找出以 \(v\) 为根的子树中出现次数最多的颜色（若多个颜色出现次数相同且均为最多，则这些颜色的编号之和）。要求输出每个节点的答案。本质是静态子树众数统计问题，需对每个子树高效统计颜色频次并维护众数信息。

2. 算法选择：

   • 树上启发式合并（DSU on Tree）：利用重链剖分思想，先处理轻儿子（不保留贡献），再处理重儿子（保留贡献），然后暴力加入轻儿子及其子树，得到当前节点的答案。若当前节点是轻儿子，则清除本次贡献（撤销）。该算法保证每个节点被加入/删除 \(O(\log n)\) 次，总时间复杂度 \(O(n \log n)\)。
   • 维护信息：用数组 cnt[color] 记录颜色出现次数；mx 记录当前子树中的最大出现次数；sum 记录所有出现次数为 mx 的颜色编号之和（即众数颜色编号之和）。

3. 关键步骤：

   • 第一次 DFS：计算每个节点的子树大小 sz[u] 和重儿子 son[u]（子树最大的儿子）。
   • 第二次 DFS（启发式合并）：
     • 对每个节点 \(u\)，先递归处理所有轻儿子，并不保留它们的统计结果（即清除贡献）。
     • 再递归处理重儿子，保留其统计结果（即重儿子的 cnt 数组直接继承）。
     • 将 \(u\) 节点自身及其所有轻子树中的节点颜色加入统计（通过 add 函数），更新 mx 和 sum。
     • 记录 ans[u] = sum。
     • 若当前节点是轻儿子（opt == 0），则撤销整个子树的贡献（sub 函数），并将 mx 和 sum 归零。
   • 辅助函数：
     • add(x, fa, son)：将节点 \(x\) 及其所有非重儿子子树（跳过 fa 和指定的 son）的颜色加入 cnt，并更新 mx 和 sum。
     • sub(x, fa)：递归减少 cnt[color[x]]，清除整个子树的颜色计数。

4. 时间/空间复杂度：

   • 时间复杂度：\(O(n \log n)\)，每个节点作为轻儿子被加入 \(O(\log n)\) 次。
   • 空间复杂度：\(O(n)\)，存储邻接表和计数数组。

5. 树上启发式合并的核心思想：

   • 重儿子优先：将重儿子的计算结果保留，避免重复计算，只暴力处理轻儿子，降低了总体复杂度。
   • 撤销机制：对于轻儿子的子树，在计算完后清除其贡献，确保不会影响其他分支的计算，保持独立性。
   • 维护众数信息：通过 cnt、mx、sum 三个变量在添加颜色时动态维护众数，无需每次重新统计。
   • 适用场景：适用于静态树上需要统计每个子树中某种信息（如众数、颜色出现次数等）的问题，是一种高效的离线算法。

【算法标签】

提高+ #树上启发式合并

【代码详解】

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

#define int long long

const int N = 100005;

int n;                                  // 节点个数
int col[N];                             // 每个节点的颜色
int sz[N], son[N];                      // 子树大小、重儿子
int fa[N], dep[N];                      // 父节点、深度（用于重链剖分，但此处未用）
int cnt[N];                             // 颜色出现次数统计
int mx;                                 // 当前子树中出现最多的次数
int sum;                                // 当前子树中众数颜色编号之和
int ans[N];                             // 每个节点的答案
vector<int> g[N];                       // 邻接表

// 将节点 x 及其轻子树（除 fa 和 son 外）的颜色加入统计
void add(int x, int fa, int son)
{
    cnt[col[x]]++;                      // 当前节点颜色计数加 1
    if (cnt[col[x]] > mx)               // 若超过当前最大次数，更新众数和
    {
        mx = cnt[col[x]];
        sum = col[x];
    }
    else if (cnt[col[x]] == mx)         // 若等于最大次数，颜色编号累加
    {
        sum += col[x];
    }

    // 递归遍历所有子节点，跳过父节点和指定的重儿子（避免重复统计）
    for (int y : g[x])
    {
        if (y != fa && y != son)
        {
            add(y, x, son);
        }
    }
}

// 从节点 x 开始，递归清除其子树中所有颜色的计数（用于撤销轻子树的影响）
void sub(int x, int fa)
{
    cnt[col[x]]--;                      // 当前节点颜色计数减 1
    for (int y : g[x])
    {
        if (y != fa)
        {
            sub(y, x);
        }
    }
}

// 第一次 DFS，计算子树大小和重儿子
void dfs1(int u, int father)
{
    fa[u] = father;
    dep[u] = dep[father] + 1;
    sz[u] = 1;

    for (int v : g[u])
    {
        if (v == father)
        {
            continue;
        }
        dfs1(v, u);
        sz[u] += sz[v];
        if (sz[son[u]] < sz[v])         // 更新重儿子（子树最大的儿子）
        {
            son[u] = v;
        }
    }
}

// 树上启发式合并（DSU on tree），opt 标记是否保留当前子树的结果
void dfs2(int x, int fa, int opt)
{
    // 先处理所有轻儿子，不保留其统计结果（opt = 0）
    for (int y : g[x])
    {
        if (y != fa && y != son[x])
        {
            dfs2(y, x, 0);
        }
    }

    // 处理重儿子，保留其统计结果（opt = 1）
    if (son[x])
    {
        dfs2(son[x], x, 1);
    }

    // 将轻儿子及其自身的颜色加入统计（此时重儿子的统计已经保留在 cnt 中）
    add(x, fa, son[x]);

    // 记录当前节点 x 的答案
    ans[x] = sum;

    // 如果不需要保留当前子树的统计（即当前是轻儿子），则清除所有颜色计数
    if (!opt)
    {
        sub(x, fa);
        sum = mx = 0;                   // 重置全局众数信息
    }
}

signed main()
{
    cin >> n;

    // 读入每个节点的颜色
    for (int i = 1; i <= n; i++)
    {
        cin >> col[i];
    }

    // 读入树边
    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    // 预处理子树大小和重儿子
    dfs1(1, 0);

    // 执行树上启发式合并
    dfs2(1, 0, 0);

    // 输出每个节点的答案
    for (int i = 1; i <= n; i++)
    {
        cout << ans[i] << " ";
    }
    cout << endl;

    return 0;
}

【运行结果】