rerooting_dp.hpp¶
两次递归 DFS 计算全根树 DP，缓存跨边贡献并共用前缀数组；合并可不交换，严格保持邻接表顺序，总计线性次状态转移。
Complexity: Time: \(O(n)\) monoid transitions on a tree. Space: \(O(n)\) values and \(O(n)\) recursion depth.
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n) values and O(n) recursion depth.
#include <cassert>
#include <vector>
namespace noya {
/// @brief Generic rerooting DP on an undirected tree; rt is a work root.
/// merge is associative with identity uni, in each adjacency list's order.
/// ae(value, from, to) moves a rooted value across an edge; av(value, vertex)
/// closes a vertex. Callbacks must return deterministic values without
/// modifying their inputs. T must be copyable, but need not be defaultable.
template <class T, class Merge, class AddEdge, class AddVertex>
std::vector<T> rerooting_dp(const std::vector<std::vector<int>> &tr,
const T &uni, Merge mer, AddEdge ae, AddVertex av,
int rt = 0) {
const int n = int(tr.size());
if (n == 0) return {};
assert(0 <= rt && rt < n);
std::size_t degree_sum = 0, max_degree = 0;
for (const auto &adj : tr) {
degree_sum += adj.size();
if (adj.size() > max_degree) max_degree = adj.size();
}
assert(degree_sum == 2 * std::size_t(n - 1));
std::vector<T> dn(n, uni), up(n, uni), ans(n, uni);
auto dfs = [&](auto &self, int u, int p) -> void {
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (v == p) continue;
self(self, v, u);
dn[v] = ae(dn[v], v, u); // Cache the child-to-parent contribution.
dn[u] = mer(dn[u], dn[v]);
}
dn[u] = av(dn[u], u);
};
dfs(dfs, rt, -1);
std::vector<T> pre(max_degree + 1, uni);
auto reroot = [&](auto &self, int u, int p) -> void {
int degree = int(tr[u].size());
for (int i = 0; i < degree; ++i) {
int v = tr[u][i];
pre[i + 1] = mer(pre[i], v == p ? up[u] : dn[v]);
}
ans[u] = av(pre[degree], u);
T suffix = uni;
for (int i = degree - 1; i >= 0; --i) {
int v = tr[u][i];
if (v != p) up[v] = ae(av(mer(pre[i], suffix), u), u, v);
suffix = mer(v == p ? up[u] : dn[v], suffix);
}
// Finish every exclusion first, so children can reuse the prefix buffer.
for (int v : tr[u]) if (v != p) self(self, v, u);
};
reroot(reroot, rt, -1);
return ans;
}
} // namespace noya
// NOYA_REROOTING_DP_HPP
#ifndef NOYA_REROOTING_DP_HPP
#define NOYA_REROOTING_DP_HPP 1
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n) values and O(n) recursion depth.
#include <cassert>
#include <vector>
namespace noya {
/// @brief Generic rerooting DP on an undirected tree; rt is a work root.
/// merge is associative with identity uni, in each adjacency list's order.
/// ae(value, from, to) moves a rooted value across an edge; av(value, vertex)
/// closes a vertex. Callbacks must return deterministic values without
/// modifying their inputs. T must be copyable, but need not be defaultable.
template <class T, class Merge, class AddEdge, class AddVertex>
std::vector<T> rerooting_dp(const std::vector<std::vector<int>> &tr,
const T &uni, Merge mer, AddEdge ae, AddVertex av,
int rt = 0) {
const int n = int(tr.size());
if (n == 0) return {};
assert(0 <= rt && rt < n);
std::size_t degree_sum = 0, max_degree = 0;
for (const auto &adj : tr) {
degree_sum += adj.size();
if (adj.size() > max_degree) max_degree = adj.size();
}
assert(degree_sum == 2 * std::size_t(n - 1));
std::vector<T> dn(n, uni), up(n, uni), ans(n, uni);
auto dfs = [&](auto &self, int u, int p) -> void {
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (v == p) continue;
self(self, v, u);
dn[v] = ae(dn[v], v, u); // Cache the child-to-parent contribution.
dn[u] = mer(dn[u], dn[v]);
}
dn[u] = av(dn[u], u);
};
dfs(dfs, rt, -1);
std::vector<T> pre(max_degree + 1, uni);
auto reroot = [&](auto &self, int u, int p) -> void {
int degree = int(tr[u].size());
for (int i = 0; i < degree; ++i) {
int v = tr[u][i];
pre[i + 1] = mer(pre[i], v == p ? up[u] : dn[v]);
}
ans[u] = av(pre[degree], u);
T suffix = uni;
for (int i = degree - 1; i >= 0; --i) {
int v = tr[u][i];
if (v != p) up[v] = ae(av(mer(pre[i], suffix), u), u, v);
suffix = mer(v == p ? up[u] : dn[v], suffix);
}
// Finish every exclusion first, so children can reuse the prefix buffer.
for (int v : tr[u]) if (v != p) self(self, v, u);
};
reroot(reroot, rt, -1);
return ans;
}
} // namespace noya
#endif // NOYA_REROOTING_DP_HPP
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n) values and O(n) recursion depth.
#include <cassert>
#include <vector>
namespace noya {
/// @brief Generic rerooting DP on an undirected tree; rt is a work root.
/// merge is associative with identity uni, in each adjacency list's order.
/// ae(value, from, to) moves a rooted value across an edge; av(value, vertex)
/// closes a vertex. Callbacks must return deterministic values without
/// modifying their inputs. T must be copyable, but need not be defaultable.
template <class T, class Merge, class AddEdge, class AddVertex>
std::vector<T> rerooting_dp(const std::vector<std::vector<int>> &tr,
const T &uni, Merge mer, AddEdge ae, AddVertex av,
int rt = 0) {
const int n = int(tr.size());
if (n == 0) return {};
assert(0 <= rt && rt < n);
std::size_t degree_sum = 0, max_degree = 0;
for (const auto &adj : tr) {
degree_sum += adj.size();
if (adj.size() > max_degree) max_degree = adj.size();
}
assert(degree_sum == 2 * std::size_t(n - 1));
std::vector<T> dn(n, uni), up(n, uni), ans(n, uni);
auto dfs = [&](auto &self, int u, int p) -> void {
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (v == p) continue;
self(self, v, u);
dn[v] = ae(dn[v], v, u); // Cache the child-to-parent contribution.
dn[u] = mer(dn[u], dn[v]);
}
dn[u] = av(dn[u], u);
};
dfs(dfs, rt, -1);
std::vector<T> pre(max_degree + 1, uni);
auto reroot = [&](auto &self, int u, int p) -> void {
int degree = int(tr[u].size());
for (int i = 0; i < degree; ++i) {
int v = tr[u][i];
pre[i + 1] = mer(pre[i], v == p ? up[u] : dn[v]);
}
ans[u] = av(pre[degree], u);
T suffix = uni;
for (int i = degree - 1; i >= 0; --i) {
int v = tr[u][i];
if (v != p) up[v] = ae(av(mer(pre[i], suffix), u), u, v);
suffix = mer(v == p ? up[u] : dn[v], suffix);
}
// Finish every exclusion first, so children can reuse the prefix buffer.
for (int v : tr[u]) if (v != p) self(self, v, u);
};
reroot(reroot, rt, -1);
return ans;
}
} // namespace noya
// NOYA_REROOTING_DP_HPP