E. Lomsat gelral - 题解

比赛与标签

比赛: Educational Codeforces Round 2 标签: data structures, dfs and similar, dsu, trees 难度: *2300

题目要我们做什么喵~

这道题给了我们一棵有根树，树根是节点1，每个节点都有一种漂亮的颜色，呐。

对于树里的每一个节点 v，我们要找到它子树里的“支配颜色” (dominating colour)。什么是支配颜色呢？如果一种颜色 c 在 v 的子树里出现的次数不比任何其他颜色少，那它就是支配颜色啦。也就是说，它是出现次数最多的颜色之一，可能会有多种颜色并列第一哦！

我们的任务就是，对每个节点 v，计算出它子树里所有支配颜色的编号之和，然后把每个节点的答案都打印出来。

举个栗子，如果节点 v 的子树里，颜色 2 出现了3次，颜色 5 出现了3次，其他颜色都少于3次，那么颜色 2 和 5 都是支配颜色，这个节点的答案就是 2 + 5 = 7 的说！

聪明的猫猫是如何思考的呐~

看到“子树查询”这种问题，一个朴素的想法就是对每个节点 v，都遍历一遍它的整个子树，用一个哈希表或者数组统计所有颜色的出现次数，找到最大频次，然后把对应颜色的编号加起来。

但是喵~ 这样做的话，对于每个节点都要重新计算一次，如果树的结构是一条长长的链，那时间复杂度就会飙到 O(N^2)，对于 N 高达 10^5 的数据量，肯定会超时的说！(；´Д｀)

所以，我们需要一种更聪明的办法，能够重复利用已经计算过的信息。这时候，一个非常适合处理这类离线子树问题的强大算法——DSU on Tree（也叫 Sack 算法）就闪亮登场啦！

DSU on Tree 的核心思想是，当我们计算完一个节点的子树信息后，尝试把这些信息“继承”给它的父节点，从而避免重复计算。具体来说，它利用了树链剖分中“重儿子”(heavy child) 的思想。

第一步：预处理，找到重儿子

我们先用一次 DFS 遍历整棵树，做两件事：

1. 计算每个节点 u 的子树大小 sz[u]。
2. 找到每个节点 u 的“重儿子” heavy[u]。重儿子就是 u 的所有孩子中，子树大小最大的那个。其他的孩子就是“轻儿子”(light child)。

第二步：DSU on Tree 主过程

我们再进行一次 DFS，这次是来计算答案的。对于当前节点 u，我们按以下顺序执行：

1. 递归处理轻儿子：先递归访问 u 的所有轻儿子。当每次递归调用结束后，清空由这个轻儿子子树带来的所有颜色统计信息。因为我们不希望一个轻儿子的信息影响到另一个轻儿子。

2. 递归处理重儿子：接着，递归访问 u 的重儿子。最关键的一步来啦！这次递归结束后，我们保留重儿子子树的所有颜色统计信息。为什么要保留呢？因为重儿子的子树是所有孩子中最大的，保留它的信息可以省去最大量的重复计算，这是整个算法高效的关键所在！

3. 合并信息：现在，我们的统计数据里已经有了重儿子子树的全部信息。我们只需要再把当前节点 u 和它所有轻儿子子树的信息“添加”进来就好啦。我们再次遍历 u 的所有轻儿子子树，把它们的颜色信息一个个加到全局的统计数据里。

4. 计算答案：当 u 和它所有子树（包括重儿子和所有轻儿子）的信息都统计完毕后，我们就可以计算出节点 u 的答案了。我们需要快速知道“出现次数最多的颜色的总和”。为了实现这一点，我们可以维护一个 max_freq 变量记录当前的最大频次，再用一个数组 total_by_freq[f] 记录所有出现频次为 f 的颜色的总和。这样，u 的答案就是 total_by_freq[max_freq] 啦！

5. 清理现场（如果需要）：最后，根据 u 是它父亲的重儿子还是轻儿子，来决定是否要清空当前的统计信息。如果 u 是一个轻儿子，那么在回溯到它父亲之前，我们就需要清空以 u 为根的整个子树的统计信息（就像步骤1里做的那样）。如果 u 是一个重儿子，我们就不做任何事，把信息留给它的父亲。

通过这种方式，每个节点的信息只会被添加和删除有限次，整个算法的复杂度就被优化到了一个非常棒的水平喵！

把思路变成代码吧，喵！

#include <iostream>
#include <vector>
#include <stack>
#include <cstring>

using namespace std;
const int N = 100005;
int n;
vector<int> adj[N]; // 邻接表存树
int color[N];       // 每个节点的颜色
int heavy[N], sz[N]; // heavy[u]是u的重儿子, sz[u]是u的子树大小
long long ans[N];   // 存储每个节点的答案

// --- 用于统计颜色的全局数据结构 ---
int cnt[N];             // cnt[c]: 颜色c出现的次数
long long total_by_freq[N]; // total_by_freq[f]: 出现次数为f的所有颜色的总和
int max_freq;           // 当前子树中的最大颜色频率

// 预处理DFS，计算子树大小并找出重儿子
void pre_dfs(int u, int p) {
    sz[u] = 1;
    heavy[u] = -1;
    int max_size = 0;
    for (int v : adj[u]) {
        if (v == p) continue;
        pre_dfs(v, u);
        sz[u] += sz[v];
        if (sz[v] > max_size) {
            max_size = sz[v];
            heavy[u] = v;
        }
    }
}

// 添加一个颜色c的贡献
void add_color(int c) {
    // 如果之前该颜色就存在，要先从旧频率的总和中减掉
    if (cnt[c] > 0) {
        total_by_freq[cnt[c]] -= c;
    }
    cnt[c]++; // 频率+1
    // 加入到新频率的总和中
    total_by_freq[cnt[c]] += c;
    // 更新全局最大频率
    if (cnt[c] > max_freq) {
        max_freq = cnt[c];
    }
}

// 移除一个颜色c的贡献
void remove_color(int c) {
    // 从当前频率的总和中减掉
    total_by_freq[cnt[c]] -= c;
    cnt[c]--; // 频率-1
    // 如果减完后频率还大于0，就加到新的频率总和里
    if (cnt[c] > 0) {
        total_by_freq[cnt[c]] += c;
    }
    // 如果最大频率的颜色总和变为0了，说明最大频率降低了
    if (max_freq > 0 && total_by_freq[max_freq] == 0) {
        max_freq--;
    }
}

// 使用迭代式DFS（栈）添加整个子树的颜色信息，避免递归爆栈
void add_subtree(int start, int parent) {
    stack<pair<int, int>> st;
    st.push({start, parent});
    while (!st.empty()) {
        auto [u, p] = st.top();
        st.pop();
        add_color(color[u]);
        for (int v : adj[u]) {
            if (v == p) continue;
            st.push({v, u});
        }
    }
}

// 使用迭代式DFS（栈）移除整个子树的颜色信息
void remove_subtree(int start, int parent) {
    stack<pair<int, int>> st;
    st.push({start, parent});
    while (!st.empty()) {
        auto [u, p] = st.top();
        st.pop();
        remove_color(color[u]);
        for (int v : adj[u]) {
            if (v == p) continue;
            st.push({v, u});
        }
    }
}

// DSU on Tree 主过程
// keep参数决定是否保留当前子树的计算结果
void dfs(int u, int p, bool keep) {
    // 1. 先递归处理所有轻儿子，并且不保留它们的结果
    for (int v : adj[u]) {
        if (v == p || v == heavy[u]) continue;
        dfs(v, u, false);
    }
    
    // 2. 递归处理重儿子，并且保留它的结果
    if (heavy[u] != -1) {
        dfs(heavy[u], u, true);
    }
    
    // 3. 将当前节点 u 和所有轻儿子子树的贡献加入
    add_color(color[u]);
    for (int v : adj[u]) {
        if (v == p || v == heavy[u]) continue;
        add_subtree(v, u);
    }
    
    // 4. 此刻，u的整个子树信息都已统计完毕，计算答案
    ans[u] = total_by_freq[max_freq];
    
    // 5. 如果当前节点是轻儿子（keep=false），则清空其子树贡献
    if (!keep) {
        remove_subtree(u, p);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> color[i];
    }
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    // 步骤一：预处理
    pre_dfs(1, 0);
    
    // 初始化全局统计变量
    max_freq = 0;
    memset(cnt, 0, sizeof(cnt));
    memset(total_by_freq, 0, sizeof(total_by_freq));
    
    // 步骤二：执行 DSU on Tree
    dfs(1, 0, false); // 从根节点开始，初始时不需要保留
    
    for (int i = 1; i <= n; i++) {
        cout << ans[i];
        if (i < n) cout << " ";
    }
    cout << endl;
    
    return 0;
}

跑得快不快，占得多不多？

• 时间复杂度: O(N log N) 的说。 为什么是这个复杂度呢？DSU on Tree 的一个神奇性质是，任何一个节点 u 到根节点的路径上，轻边的数量不会超过 log N 条。当我们在 dfs 函数中处理一个节点时，只有当它是一个轻儿子的子树的一部分时，它的信息才会被重复地 add 和 remove。由于轻边路径的限制，每个节点最多被作为轻子树的部分遍历 log N 次。所以总的时间复杂度就是 O(N log N) 啦，非常高效！

• 空间复杂度: O(N) 的说。 我们需要邻接表 adj、颜色数组 color、子树大小 sz 等等，这些都是 O(N) 的。用来统计的 cnt 数组和 total_by_freq 数组大小也和 N 相关，所以总空间复杂度是 O(N)，完全没问题！

猫猫的小课堂时间~

这道题真是一道非常经典的 DSU on Tree 模板题呢，喵~ 通过它，我们可以学到很多东西：

1. DSU on Tree (Sack) 算法: 这是解决一类离线子树询问问题的通用利器。它的核心就是“重儿子启发式合并”，通过保留重儿子的信息来避免大量重复计算。如果你以后遇到类似“统计子树中XXX”且没有修改操作的问题，一定要想想它哦！

2. 优雅地统计答案: 题解中的 max_freq 和 total_by_freq 数组配合得天衣无缝！只用 O(1) 的时间就能更新颜色计数和查询答案，这个小技巧非常值得学习。它把“找到最大频次”和“求和”这两个步骤完美地融合在了数据更新的过程中。

3. 迭代式DFS: 代码中的 add_subtree 和 remove_subtree 使用了栈来实现DFS，这是一个很好的编程习惯。在处理深度可能很大的树时，可以有效防止递归层数过多导致的栈溢出（Stack Overflow）问题。

总之，只要理解了 DSU on Tree 的思想，这道题就会变得非常清晰。希望这篇题解能帮助大家掌握这个强大的算法，以后遇到类似的题目也能轻松解决啦！加油喵~ (ฅ'ω'ฅ)