D. Tree and Queries - 题解

比赛与标签

比赛: Codeforces Round 221 (Div. 1) 标签: data structures, dfs and similar, trees 难度: *2400

题目大意喵~

主人，你好呀！这次的任务是关于一棵可爱的、五颜六色的树哦~ ฅ^•ﻌ•^ฅ

我们有一棵 n 个节点的有根树，树根是节点 1。每个节点 v 都有一个自己的颜色 c_v。

接下来会有 m 个询问。每个询问会给我们一个节点 v_j 和一个数字 k_j。我们需要回答的是：在以 v_j 为根的子树里，有多少种不同的颜色，它们的出现次数至少是 k_j 次？

简单来说，就是对每个查询 (v, k)，找出 v 的子树里所有节点，统计每种颜色的出现次数，然后数一数有多少种颜色的出现次数大于等于 k。

解题思路喵！DSU on Tree 的奇妙魔法！

看到是树上的子树查询问题，我们的第一反应可能是对每个查询 (v, k)，都做一次 DFS 或 BFS，遍历 v 的整个子树，用一个 map 或者哈希表统计颜色数量，然后再回答查询。

但是呀，n 和 m 都高达 10^5，如果每个查询都暴力遍历一遍，当树形成一条链的时候，最坏情况下复杂度会飙到 O(n*m)，这肯定会超时的说！(T_T)

所以，我们需要更聪明的办法！这种子树问题，而且不带修改，很适合离线处理。我们可以把所有查询挂在对应的节点 v 上，然后在一次 DFS 遍历中解决所有问题。

这时候，DSU on Tree（也被称为 Sack 或树上启发式合并）就闪亮登场啦！这是一种专门优美地解决这类子树统计问题的黑魔法，喵~

它的核心思想是：利用重链剖分中的“重儿子”思想，来减少重复计算。

对于一个节点 u，它的儿子中，子树最大的那个被称为重儿子 (big child)，其他的都是轻儿子 (light child)。

DSU on Tree 的魔法步骤如下：

1. 预处理 dfs_sz: 先进行一次 DFS，计算出每个节点的子树大小 sz[u]，并找出每个节点的重儿子 big_child[u]。
2. 主过程 dfs_solve(u, keep): 这是第二次 DFS，也是算法的核心！
   • 首先，递归处理 u 的所有轻儿子，并且在处理完每个轻儿子后，清除它们对全局统计数据（比如颜色计数器）造成的影响。就像小猫咪来玩了一下就跑了，不留下痕迹~
   • 然后，递归处理 u 的重儿子。最关键的一步来啦：处理完重儿子后，我们保留它对全局统计数据的影响！因为重儿子的子树最大，保留它的信息可以省下最多的重复计算，这可是非常划算的买卖呢！
   • 现在，我们的统计数据里已经包含了 u 的重儿子子树的所有信息。接下来，我们把 u 节点本身，以及 u 所有轻儿子子树的信息，再一个个添加进去。
   • 当当当当！ 在这一刻，我们的全局统计数据正好完整地包含了整个 u 子树的信息！太神奇了对不对？这时候我们就可以愉快地回答所有挂在节点 u 上的查询了。
   • 最后，根据 keep 参数决定是否要“打扫战场”。如果 keep 是 false（意味着 u 是它父亲的轻儿子），我们就需要清除 u 整个子树造成的影响，为回溯到它的兄弟节点做准备。如果 keep 是 true（u 是它父亲的重儿子），那我们就啥也不用干，把这些宝贵的信息留给它的父亲用，喵~

如何高效回答查询 "有多少种颜色出现次数 >= k" 呢？

我们用一个数组 color_counts[c] 来记录颜色 c 在当前子树中出现了多少次。但如果每次查询都遍历一遍 color_counts，还是太慢了。

这里就需要另一个好朋友——树状数组 (Fenwick Tree) 啦！

我们可以用树状数组来维护频率的频率。听起来有点绕？别怕，是这样的： 我们让树状数组 ft 的第 i 个位置 ft[i] 存储 “当前出现次数恰好为 i 的颜色有多少种”。

• 当某个颜色 c 的出现次数从 x 变为 x+1 时：
  1. 出现次数为 x 的颜色种类数减一，所以 ft_update(x, -1)。
  2. 出现次数为 x+1 的颜色种类数加一，所以 ft_update(x+1, 1)。
• 当我们要查询出现次数大于等于 k 的颜色有多少种时，我们只需要计算 ft[k] + ft[k+1] + ... + ft[n] 的和。这正是树状数组最擅长的区间查询！用 ft_query(n) - ft_query(k-1) 就可以在 O(log n) 的时间内得到答案啦！

把 DSU on Tree 和树状数组结合起来，我们就有了一个超级强大的解决方案！

代码实现喵~

#include <iostream>
#include <vector>
#include <algorithm>

const int MAX_N_M = 100005;
const int MAX_COLORS = 100005;

int n, m;
int colors[MAX_N_M];
std::vector<int> adj[MAX_N_M];
// queries[v] 存储所有关于节点 v 的查询 {k, query_index}
std::vector<std::pair<int, int>> queries[MAX_N_M]; 
int ans[MAX_N_M];

// DSU on Tree 需要的辅助数组
int sz[MAX_N_M];      // sz[u]：以 u 为根的子树大小
int big_child[MAX_N_M]; // big_child[u]：u 的重儿子

// 核心数据结构
int color_counts[MAX_COLORS]; // color_counts[c]：颜色 c 在当前子树中的出现次数
int ft[MAX_N_M + 1];          // 树状数组，ft[i] 存储出现次数为 i 的颜色有多少种

// --- 树状数组操作 ---
void ft_update(int idx, int delta) {
    for (; idx <= n; idx += idx & -idx) {
        ft[idx] += delta;
    }
}

int ft_query(int idx) {
    int sum = 0;
    for (; idx > 0; idx -= idx & -idx) {
        sum += ft[idx];
    }
    return sum;
}

// 查询出现次数在 [l, r] 区间内的颜色种类数
int ft_query_range(int l, int r) {
    if (l > r) return 0;
    return ft_query(r) - ft_query(l - 1);
}

// --- 颜色计数更新 ---
// 更新颜色 c 的计数，val 可以是 +1 (添加) 或 -1 (移除)
void update_color(int c, int val) {
    // 1. 从树状数组中移除旧的计数贡献
    if (color_counts[c] > 0) {
        ft_update(color_counts[c], -1);
    }
    // 2. 更新颜色计数
    color_counts[c] += val;
    // 3. 在树状数组中添加新的计数贡献
    if (color_counts[c] > 0) {
        ft_update(color_counts[c], 1);
    }
}

// --- DSU on Tree 核心 DFS ---

// 第一次DFS：计算子树大小和确定重儿子
void dfs_sz(int u, int p) {
    sz[u] = 1;
    int max_sz = 0;
    big_child[u] = -1;
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs_sz(v, u);
        sz[u] += sz[v];
        if (sz[v] > max_sz) {
            max_sz = sz[v];
            big_child[u] = v;
        }
    }
}

// 递归地将 u 子树的所有节点颜色信息添加到全局计数中
void add_subtree(int u, int p) {
    update_color(colors[u], 1);
    for (int v : adj[u]) {
        if (v == p) continue;
        add_subtree(v, u);
    }
}

// 递归地将 u 子树的所有节点颜色信息从全局计数中移除
void remove_subtree(int u, int p) {
    update_color(colors[u], -1);
    for (int v : adj[u]) {
        if (v == p) continue;
        remove_subtree(v, u);
    }
}

// 第二次DFS：实现 DSU on Tree 算法，解决查询
void dfs_solve(int u, int p, bool keep) {
    // 1. 递归处理所有轻儿子，处理完后清除它们的数据 (keep=false)
    for (int v : adj[u]) {
        if (v != p && v != big_child[u]) {
            dfs_solve(v, u, false);
        }
    }

    // 2. 递归处理重儿子，处理完后保留它的数据 (keep=true)
    if (big_child[u] != -1) {
        dfs_solve(big_child[u], u, true);
    }

    // 3. 将当前节点 u 和所有轻儿子子树的信息添加进来
    update_color(colors[u], 1);
    for (int v : adj[u]) {
        if (v != p && v != big_child[u]) {
            add_subtree(v, u);
        }
    }

    // 4. 此刻，数据结构中包含了 u 的完整子树信息，回答所有关于 u 的查询
    for (auto const& [k, q_idx] : queries[u]) {
        // 查询出现次数 >= k 的颜色数量
        ans[q_idx] = ft_query_range(k, n);
    }

    // 5. 如果 u 是轻儿子 (keep=false)，清除整个 u 子树的信息
    if (!keep) {
        remove_subtree(u, p);
    }
}

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    std::cin >> n >> m;
    for (int i = 1; i <= n; ++i) {
        std::cin >> colors[i];
    }
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    // 离线存储查询
    for (int i = 0; i < m; ++i) {
        int v, k;
        std::cin >> v >> k;
        queries[v].push_back({k, i});
    }

    // 启动 DSU on Tree 魔法！
    dfs_sz(1, 0);
    dfs_solve(1, 0, false);

    // 按顺序输出答案
    for (int i = 0; i < m; ++i) {
        std::cout << ans[i] << "\n";
    }

    return 0;
}

复杂度分析🐾

• 时间复杂度: O(N log^2 N + M log N) 的说

  • dfs_sz 预处理是 O(N) 的。
  • dfs_solve 是主角。DSU on Tree 的一个关键性质是：从根到任意节点的路径上，轻边的数量不超过 O(log N) 条。这意味着每个节点 v 的信息，只会在它作为某个祖先的轻子树成员时被暴力添加和删除。这样的祖先最多有 O(log N) 个。
  • 所以，每个节点被 add_subtree 和 remove_subtree 访问的总次数是 O(log N) 次。
  • 每次访问（调用 update_color）会操作树状数组，花费 O(log N) 时间。
  • 因此，所有节点更新的总时间是 O(N log N * log N) = O(N log^2 N)。
  • 回答 m 个查询需要 O(M log N) 的时间。
  • 总时间复杂度就是 O(N log^2 N + M log N) 啦！非常高效！

• 空间复杂度: O(N + M) 的说

  • 邻接表 adj 需要 O(N) 空间。
  • colors, sz, big_child 等辅助数组都是 O(N) 的。
  • 存储所有查询的 queries 数组需要 O(M) 空间。
  • color_counts 和 ft 树状数组最大需要 O(N) 或 O(颜色种类数) 的空间。
  • DFS 的递归栈深度最多为 O(N)。
  • 所以总的空间复杂度是 O(N + M) 哦。

知识点与总结，喵~

这道题真是一道非常经典的模板题，把好几种算法思想巧妙地结合在了一起，做完之后收获满满呢！

1. DSU on Tree (Sack / 树上启发式合并): 这是解决本题的核心科技！它是一种处理离线树上子树问题的强大技巧，通过“保留重儿子，暴力轻儿子”的策略，将暴力算法的复杂度优化到了一个非常优秀的级别。一定要掌握它的思想和代码模板哦！

2. 离线处理: 当题目查询不涉及修改，并且可以一次性读入时，离线处理是一个非常好的思路。把查询挂在对应的节点上，在一次遍历中统一解决，可以避免大量重复计算。

3. 树状数组 (Fenwick Tree): 在这道题里，树状数组扮演了高效计数器的角色。它巧妙地维护了“频率的频率”，将一个看似复杂的查询（“有多少种颜色出现次数>=k”）转化为了一个简单的区间求和问题，把每次查询的复杂度从 O(N) 降到了 O(log N)。

总之，当你遇到树上子树的统计问题，而且没有在线修改时，一定要想想 DSU on Tree 这个好朋友呀！它能帮你解决很多看起来很棘手的问题，喵~ 希望这篇题解能帮助到你，一起加油，成为更厉害的算法大师吧！(๑•̀ㅂ•́)و✧