dsu_on_tree.hpp¶
按重儿子保留统计状态，离线计算每个子树的答案；适合能逐点加入/删除贡献的子树询问。
Complexity: Time: O(n log n) callback invocations on a tree. Space: O(n).
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n log n) callback invocations on a tree.
/// Space: O(n).
#include "noya/graph/heavy_light_decomposition.hpp"
namespace noya {
/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
auto &&del) {
hld hl1(g);
auto dfs = [&](auto &dfs, int u, bool so1) -> void {
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
dfs(dfs, v, false);
}
}
if (hl1.son[u] != -1)
dfs(dfs, hl1.son[u], true);
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
auto [l, r] = hl1.subtree(v);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
qry(u, x);
}
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
add(x);
}
}
}
qry(u, u);
add(u);
if (!so1) {
auto [l, r] = hl1.subtree(u);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
del(x);
}
}
};
dfs(dfs, 0, 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 on a tree.
/// Space: O(n).
#include "noya/graph/heavy_light_decomposition.hpp"
namespace noya {
/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
auto &&del) {
hld hl1(g);
auto dfs = [&](auto &dfs, int u, bool so1) -> void {
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
dfs(dfs, v, false);
}
}
if (hl1.son[u] != -1)
dfs(dfs, hl1.son[u], true);
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
auto [l, r] = hl1.subtree(v);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
qry(u, x);
}
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
add(x);
}
}
}
qry(u, u);
add(u);
if (!so1) {
auto [l, r] = hl1.subtree(u);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
del(x);
}
}
};
dfs(dfs, 0, true);
}
} // namespace noya
#endif // NOYA_DSU_ON_TREE_HPP
/// @complexity Time: O(n log n) callback invocations on a tree.
/// Space: O(n).
/// @complexity Time: O(n) build, O(log n) LCA/path jumps/decomposition,
/// O(1) subtree interval. Space: O(n).
/// @complexity Time: O(n) build; common queries add O(1) work to their LCA calls,
/// except path(), which takes O(path length). Space: O(n).
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
namespace noya {
/// @brief Iterative rooted-tree preprocessing shared by RMQ LCA and HLD.
/// The graph is a connected tree, with integer endpoints or (endpoint, weight)
/// pairs. Root depth is zero, root parent is -1, and subtree ranges are [l, r).
/// No adjacency-list copy, RMQ table, jump table, or heavy paths are retained.
struct tree_preprocessing {
int n = 0;
int rt = -1;
std::vector<int> fa, d, dfn, siz, tl;
tree_preprocessing() = default;
template <class Edge>
explicit tree_preprocessing(const std::vector<std::vector<Edge>> &g,
int root = 0) {
build(g, root);
}
template <class Edge> static int endpoint(const Edge &edge) {
if constexpr (std::is_integral_v<Edge>) {
return int(edge);
} else {
return int(std::get<0>(edge));
}
}
template <class Edge>
void build(const std::vector<std::vector<Edge>> &g, int root = 0) {
n = int(g.size());
rt = n == 0 ? -1 : root;
fa.assign(n, -1);
d.assign(n, 0);
dfn.assign(n, -1);
siz.assign(n, 1);
tl.clear();
tl.reserve(n);
if (n == 0) {
return;
}
assert(0 <= root && root < n);
std::vector<int> stack{root};
while (!stack.empty()) {
int u = stack.back();
stack.pop_back();
assert(dfn[u] == -1);
dfn[u] = int(tl.size());
tl.push_back(u);
for (auto it = g[u].rbegin(); it != g[u].rend(); ++it) {
int v = endpoint(*it);
assert(0 <= v && v < n);
if (v != fa[u]) {
fa[v] = u;
d[v] = d[u] + 1;
stack.push_back(v);
}
}
}
assert(int(tl.size()) == n);
for (auto it = tl.rbegin(); it != tl.rend(); ++it) {
if (fa[*it] != -1) {
siz[fa[*it]] += siz[*it];
}
}
}
void build(const std::vector<std::vector<int>> &g = {}, int root = 0) {
build<int>(g, root);
}
bool is_subtree(int a, int b) const {
assert(0 <= a && a < n && 0 <= b && b < n);
return dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b];
}
std::array<int, 2> subtree(int u) const {
assert(0 <= u && u < n);
return {dfn[u], dfn[u] + siz[u]};
}
};
/// @brief Shared tree queries with compile-time dispatch to Derived::lca().
template <class Derived> struct tree_queries : tree_preprocessing {
int rooted_lca(int a, int b, int root) const {
const auto &self = static_cast<const Derived &>(*this);
return self.lca(a, b) ^ self.lca(a, root) ^ self.lca(b, root);
}
std::int64_t distance(int a, int b) const {
const auto &self = static_cast<const Derived &>(*this);
int c = self.lca(a, b);
return std::int64_t(d[a]) + d[b] - 2 * std::int64_t(d[c]);
}
std::pair<int, int> intersection(int a, int b, int c, int e) const {
const auto &self = static_cast<const Derived &>(*this);
int ab = self.lca(a, b), ac = self.lca(a, c), ae = self.lca(a, e);
int bc = self.lca(b, c), be = self.lca(b, e), ce = self.lca(c, e);
int x = ab ^ ac ^ bc;
int y = ab ^ ae ^ be;
if (x != y) {
return {x, y};
}
int z = ac ^ ae ^ ce;
return x == z ? std::pair{x, x} : std::pair{-1, -1};
}
std::pair<int, int> intersection(std::pair<int, int> a,
std::pair<int, int> b) const {
return intersection(a.first, a.second, b.first, b.second);
}
std::vector<int> path(int a, int b) const {
const auto &self = static_cast<const Derived &>(*this);
int c = self.lca(a, b);
std::vector<int> left, right;
while (a != c) {
left.push_back(a);
a = fa[a];
}
while (b != c) {
right.push_back(b);
b = fa[b];
}
left.push_back(c);
left.insert(left.end(), right.rbegin(), right.rend());
return left;
}
};
} // namespace noya
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
/// @brief Iterative heavy-light decomposition for paths, jumps, and subtrees.
/// Root depth is zero. No graph copy, sparse table, or binary jump table is stored.
struct hld : tree_queries<hld> {
std::vector<int> son, top;
explicit hld(const std::vector<std::vector<int>> &g = {}, int root = 0) {
build(g, root);
}
void build(const std::vector<std::vector<int>> &g = {}, int root = 0) {
tree_preprocessing::build(g, root);
son.assign(n, -1);
top.assign(n, -1);
for (int u : tl) {
for (int v : g[u]) {
if (fa[v] == u && (son[u] == -1 || siz[v] > siz[son[u]])) {
son[u] = v;
}
}
}
tl.clear();
if (n == 0) {
return;
}
std::vector<int> stack{root};
while (!stack.empty()) {
int first = stack.back();
stack.pop_back();
for (int u = first; u != -1; u = son[u]) {
top[u] = first;
dfn[u] = int(tl.size());
tl.push_back(u);
for (auto it = g[u].rbegin(); it != g[u].rend(); ++it) {
int v = *it;
if (fa[v] == u && v != son[u]) {
stack.push_back(v);
}
}
}
}
}
/// @brief Return the k-th ancestor, or -1 for a negative/out-of-range k.
int kth_ancestor(int u, int k) const {
assert(0 <= u && u < n);
if (k < 0 || k > d[u]) {
return -1;
}
int target_depth = d[u] - k;
while (d[top[u]] > target_depth) {
u = fa[top[u]];
}
return tl[dfn[u] - (d[u] - target_depth)];
}
/// @brief Return the zero-indexed k-th vertex on a -> b, or -1 out of range.
int kth_on_path(int a, int b, int k) const {
int c = lca(a, b);
int left = d[a] - d[c], right = d[b] - d[c];
if (k < 0 || k > left + right) {
return -1;
}
return k <= left ? kth_ancestor(a, k) : kth_ancestor(b, left + right - k);
}
int lca(int x, int y) const {
assert(0 <= x && x < n && 0 <= y && y < n);
while (top[x] != top[y]) {
if (d[top[x]] < d[top[y]]) {
std::swap(x, y);
}
x = fa[top[x]];
}
return d[x] < d[y] ? x : y;
}
/// @brief Decompose x -> y into (dfn_l, dfn_r, forward) half-open intervals,
/// in path order. forward=false means traverse the interval in reverse.
std::vector<std::tuple<int, int, bool>> chain(int x, int y) const {
assert(0 <= x && x < n && 0 <= y && y < n);
std::vector<std::tuple<int, int, bool>> left, right;
while (top[x] != top[y]) {
if (d[top[x]] > d[top[y]]) {
left.emplace_back(dfn[top[x]], dfn[x] + 1, false);
x = fa[top[x]];
} else {
right.emplace_back(dfn[top[y]], dfn[y] + 1, true);
y = fa[top[y]];
}
}
if (dfn[y] < dfn[x]) {
left.emplace_back(dfn[y], dfn[x] + 1, false);
} else {
right.emplace_back(dfn[x], dfn[y] + 1, true);
}
std::reverse(right.begin(), right.end());
left.insert(left.end(), right.begin(), right.end());
return left;
}
};
} // namespace noya
namespace noya {
/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
auto &&del) {
hld hl1(g);
auto dfs = [&](auto &dfs, int u, bool so1) -> void {
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
dfs(dfs, v, false);
}
}
if (hl1.son[u] != -1)
dfs(dfs, hl1.son[u], true);
for (auto v : g[u]) {
if (v != hl1.fa[u] && v != hl1.son[u]) {
auto [l, r] = hl1.subtree(v);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
qry(u, x);
}
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
add(x);
}
}
}
qry(u, u);
add(u);
if (!so1) {
auto [l, r] = hl1.subtree(u);
for (int i = l; i < r; i++) {
int x = hl1.tl[i];
del(x);
}
}
};
dfs(dfs, 0, true);
}
} // namespace noya
// NOYA_DSU_ON_TREE_HPP