Skip to content

dsu_on_tree.hpp

SECTIONGraph INCLUDEnoya/graph/dsu_on_tree.hpp

先处理轻子树并清空，再保留重子树、合并轻子树；用加入、查询、删除三个回调统计子树或跨子树贡献。

Complexity: Time: \(O(n \log n)\) callback invocations; Space: \(O(n)\).

跳到代码 · GitHub ↗

Implementation

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(n log n) callback invocations; Space: O(n).

#include <vector>

namespace noya {

/// @brief DSU on a tree rooted at 0; empty input is a no-op, g is unchanged.
/// add(x) inserts x, qry(u,x) queries its contribution at u, del(x) removes x.
/// Query a whole light subtree before adding it; qry(u,u) precedes add(u).
/// Keep the root's final state. Recursion depth is O(tree height).
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  if (g.empty()) return;
  std::vector<int> son(g.size(), -1);
  auto init = [&](auto &self, int u, int p) -> int {
    int sz = 1, best = 0;
    for (int v : g[u]) if (v != p) {
      int s = self(self, v, u);
      sz += s;
      if (s > best) best = s, son[u] = v;
    }
    return sz;
  };
  init(init, 0, -1);

  // Visit a subtree in heavy-first order, preserving callback order.
  auto walk = [&](auto &self, int u, int p, auto &f) -> void {
    f(u);
    if (son[u] != -1) self(self, son[u], u, f);
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, f);
  };
  auto dfs = [&](auto &self, int u, int p, bool keep) -> void {
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, false);
    if (son[u] != -1) self(self, son[u], u, true);
    auto query = [&](int x) { qry(u, x); };
    for (int v : g[u]) if (v != p && v != son[u]) {
      walk(walk, v, u, query);
      walk(walk, v, u, add);
    }
    qry(u, u);
    add(u);
    if (!keep) walk(walk, u, p, del);
  };
  dfs(dfs, 0, -1, true);
}

} // namespace noya

 // NOYA_DSU_ON_TREE_HPP

#ifndef NOYA_DSU_ON_TREE_HPP
#define NOYA_DSU_ON_TREE_HPP 1

/// @complexity Time: O(n log n) callback invocations; Space: O(n).

#include <vector>

namespace noya {

/// @brief DSU on a tree rooted at 0; empty input is a no-op, g is unchanged.
/// add(x) inserts x, qry(u,x) queries its contribution at u, del(x) removes x.
/// Query a whole light subtree before adding it; qry(u,u) precedes add(u).
/// Keep the root's final state. Recursion depth is O(tree height).
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  if (g.empty()) return;
  std::vector<int> son(g.size(), -1);
  auto init = [&](auto &self, int u, int p) -> int {
    int sz = 1, best = 0;
    for (int v : g[u]) if (v != p) {
      int s = self(self, v, u);
      sz += s;
      if (s > best) best = s, son[u] = v;
    }
    return sz;
  };
  init(init, 0, -1);

  // Visit a subtree in heavy-first order, preserving callback order.
  auto walk = [&](auto &self, int u, int p, auto &f) -> void {
    f(u);
    if (son[u] != -1) self(self, son[u], u, f);
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, f);
  };
  auto dfs = [&](auto &self, int u, int p, bool keep) -> void {
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, false);
    if (son[u] != -1) self(self, son[u], u, true);
    auto query = [&](int x) { qry(u, x); };
    for (int v : g[u]) if (v != p && v != son[u]) {
      walk(walk, v, u, query);
      walk(walk, v, u, add);
    }
    qry(u, u);
    add(u);
    if (!keep) walk(walk, u, p, del);
  };
  dfs(dfs, 0, -1, true);
}

} // namespace noya

#endif // NOYA_DSU_ON_TREE_HPP

/// @complexity Time: O(n log n) callback invocations; Space: O(n).

#include <vector>

namespace noya {

/// @brief DSU on a tree rooted at 0; empty input is a no-op, g is unchanged.
/// add(x) inserts x, qry(u,x) queries its contribution at u, del(x) removes x.
/// Query a whole light subtree before adding it; qry(u,u) precedes add(u).
/// Keep the root's final state. Recursion depth is O(tree height).
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  if (g.empty()) return;
  std::vector<int> son(g.size(), -1);
  auto init = [&](auto &self, int u, int p) -> int {
    int sz = 1, best = 0;
    for (int v : g[u]) if (v != p) {
      int s = self(self, v, u);
      sz += s;
      if (s > best) best = s, son[u] = v;
    }
    return sz;
  };
  init(init, 0, -1);

  // Visit a subtree in heavy-first order, preserving callback order.
  auto walk = [&](auto &self, int u, int p, auto &f) -> void {
    f(u);
    if (son[u] != -1) self(self, son[u], u, f);
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, f);
  };
  auto dfs = [&](auto &self, int u, int p, bool keep) -> void {
    for (int v : g[u])
      if (v != p && v != son[u]) self(self, v, u, false);
    if (son[u] != -1) self(self, son[u], u, true);
    auto query = [&](int x) { qry(u, x); };
    for (int v : g[u]) if (v != p && v != son[u]) {
      walk(walk, v, u, query);
      walk(walk, v, u, add);
    }
    qry(u, u);
    add(u);
    if (!keep) walk(walk, u, p, del);
  };
  dfs(dfs, 0, -1, true);
}

} // namespace noya

 // NOYA_DSU_ON_TREE_HPP