树上启发式合并⚓︎

树上启发式合并（ $\text{DSU on Tree}$ ）是一种用于处理树形数据结构上的查询和更新操作的技术。它结合了并查集（ $\text{Disjoint Set Union, DSU}$ ）和树的深度优先搜索（ $\text{DFS}$ ）方法，能够高效地解决一些复杂的问题，如子树查询、路径查询等。

在并查集中常见的启发式合并策略包括按秩合并和按大小合并。树上启发式合并则利用了树的结构特点，通过区分不同子树的大小，优先合并较小的子树，从而减少不必要的计算和内存开销。

树上启发式合并的特征：

没有修改操作
可以通过遍历子树建立信息统计，得到所有查询的答案

树上启发式合并使用深度优先搜索（ $\text{DFS}$ ）遍历树的节点，附加一个 $keep$ 标记表示是否保留子树对信息的贡献，并在遍历过程中维护一个数据结构来存储当前子树的信息。具体步骤如下：

先遍历所有轻儿子的子树，遍历结束时，消除对信息的贡献。即： $dfs(轻儿子, false)$
再遍历唯一重儿子的子树，遍历结束时，保留对信息的贡献。即： $dfs(重儿子, true)$
考察单个节点 $u$ ，对信息进行贡献
再遍历所有轻儿子的子树上面的每个节点，都重新对信息进行贡献
得到子树 $u$ 的答案
如果 $keep == false$ ，消除子树 $u$ 的贡献；如果 $keep == true$ ，保留子树 $u$ 的贡献

时间复杂度： $O(N \log N)$

树上数颜色

给定一棵有 $N$ 个节点的树，每个节点有一个颜色。对于每个查询，求以节点 $u$ 为根的子树中有多少种不同的颜色。

C++

#include <cstdint>
#include <iostream>
#include <vector>
using namespace std;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
  int64_t n;
  cin >> n;

  vector<vector<int64_t>> tree(n + 1);
  for (int64_t i = 1; i < n; ++i) {
    int64_t u, v;
    cin >> u >> v;
    tree[u].emplace_back(v);
    tree[v].emplace_back(u);
  }

  vector<int64_t> color(n + 1);
  for (int64_t i = 1; i <= n; ++i) { cin >> color[i]; }

  int64_t r = 1;  // 根节点
  // 父节点, 深度, 子树大小, 重儿子
  vector<int64_t> parent(n + 1), depth(n + 1), size(n + 1), heavy_son(n + 1, -1);
  auto dfs = [&](auto &self, int64_t u, int64_t from) -> void {
    parent[u]        = from;
    size[u]          = 1;
    int64_t max_size = 0;
    for (int64_t v : tree[u]) {
      if (v != from) {
        depth[v] = depth[u] + 1;
        self(self, v, u);
        size[u] += size[v];
        if (size[v] > max_size) {
          max_size     = size[v];
          heavy_son[u] = v;
        }
      }
    }
  };
  dfs(dfs, r, -1);  // 计算父节点, 深度, 子树大小, 重儿子

  // 颜色计数器, 答案数组
  vector<int64_t> counter(n + 1), answer(n + 1);
  int64_t unique_colors = 0;

  // 树上启发式合并
  auto effect = [&](auto &&self, int64_t u) -> void {
    if (++counter[color[u]] == 1) { unique_colors++; }
    for (int64_t v : tree[u]) {
      if (v != parent[u]) { self(self, v); }
    }
  };

  auto deffect = [&](auto &&self, int64_t u) -> void {
    if (--counter[color[u]] == 0) { unique_colors--; }
    for (int64_t v : tree[u]) {
      if (v != parent[u]) { self(self, v); }
    }
  };

  auto dsu = [&](auto &self, int64_t u, bool keep) -> void {
    // 处理轻儿子
    for (int64_t v : tree[u]) {
      if (v != parent[u] && v != heavy_son[u]) { self(self, v, false); }
    }
    // 处理重儿子
    if (heavy_son[u] != -1) { self(self, heavy_son[u], true); }
    // 当前节点信息
    if (++counter[color[u]] == 1) { unique_colors++; }
    // 合并轻儿子的信息
    for (int64_t v : tree[u]) {
      if (v != parent[u] && v != heavy_son[u]) { effect(effect, v); }
    }
    answer[u] = unique_colors;
    // 如果不保留信息, 则清空当前子树的信息
    if (!keep) { deffect(deffect, u); }
  };
  dsu(dsu, r, false);

  int64_t m;
  cin >> m;
  for (int64_t i = 0; i < m; ++i) {
    int64_t u;
    cin >> u;
    cout << answer[u] << "\n";
  }
}

树上启发式合并⚓︎

评论