rerooting_dp.hpp¶
在线性时间内计算以每个顶点为根的树 DP；适合答案由相邻子树贡献可结合得到的全根查询。
Complexity: Time: O(n) monoid transitions on a tree. Space: O(n).
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n).
#include <cassert>
#include <vector>
namespace noya {
/// @brief Generic rerooting DP on a tree.
/// merge combines neighbor contributions, ae(value, from, to) moves a
/// rooted value across an edge, and av(value, vertex) closes a vertex.
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);
int ds = 0;
for (const auto &adj : tr) {
ds += int(adj.size());
}
assert(ds == 2 * (n - 1));
std::vector<int> fa(n, -1);
std::vector<int> ord = {rt};
fa[rt] = rt;
for (int i = 0; i < int(ord.size()); i++) {
int u = ord[i];
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (fa[v] == -1) {
fa[v] = u;
ord.push_back(v);
}
}
}
assert(int(ord.size()) == n);
std::vector<T> dn(n, uni);
for (int i = n - 1; i >= 0; i--) {
int u = ord[i];
T val = uni;
for (int v : tr[u]) {
if (fa[v] == u) {
val = mer(val, ae(dn[v], v, u));
}
}
dn[u] = av(val, u);
}
std::vector<T> up(n, uni);
std::vector<T> ans(n, uni);
for (int u : ord) {
int deg = int(tr[u].size());
std::vector<T> par(deg, uni);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
par[i] = v == fa[u] ? up[u] : ae(dn[v], v, u);
}
std::vector<T> pre(deg + 1, uni);
std::vector<T> suf(deg + 1, uni);
for (int i = 0; i < deg; i++) {
pre[i + 1] = mer(pre[i], par[i]);
}
for (int i = deg - 1; i >= 0; i--) {
suf[i] = mer(par[i], suf[i + 1]);
}
ans[u] = av(pre[deg], u);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
if (fa[v] == u) {
T wc = mer(pre[i], suf[i + 1]);
up[v] = ae(av(wc, u), u, v);
}
}
}
return ans;
}
} // namespace noya
#ifndef NOYA_REROOTING_DP_HPP
#define NOYA_REROOTING_DP_HPP 1
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n).
#include <cassert>
#include <vector>
namespace noya {
/// @brief Generic rerooting DP on a tree.
/// merge combines neighbor contributions, ae(value, from, to) moves a
/// rooted value across an edge, and av(value, vertex) closes a vertex.
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);
int ds = 0;
for (const auto &adj : tr) {
ds += int(adj.size());
}
assert(ds == 2 * (n - 1));
std::vector<int> fa(n, -1);
std::vector<int> ord = {rt};
fa[rt] = rt;
for (int i = 0; i < int(ord.size()); i++) {
int u = ord[i];
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (fa[v] == -1) {
fa[v] = u;
ord.push_back(v);
}
}
}
assert(int(ord.size()) == n);
std::vector<T> dn(n, uni);
for (int i = n - 1; i >= 0; i--) {
int u = ord[i];
T val = uni;
for (int v : tr[u]) {
if (fa[v] == u) {
val = mer(val, ae(dn[v], v, u));
}
}
dn[u] = av(val, u);
}
std::vector<T> up(n, uni);
std::vector<T> ans(n, uni);
for (int u : ord) {
int deg = int(tr[u].size());
std::vector<T> par(deg, uni);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
par[i] = v == fa[u] ? up[u] : ae(dn[v], v, u);
}
std::vector<T> pre(deg + 1, uni);
std::vector<T> suf(deg + 1, uni);
for (int i = 0; i < deg; i++) {
pre[i + 1] = mer(pre[i], par[i]);
}
for (int i = deg - 1; i >= 0; i--) {
suf[i] = mer(par[i], suf[i + 1]);
}
ans[u] = av(pre[deg], u);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
if (fa[v] == u) {
T wc = mer(pre[i], suf[i + 1]);
up[v] = ae(av(wc, u), u, v);
}
}
}
return ans;
}
} // namespace noya
#endif // NOYA_REROOTING_DP_HPP
#include <cassert>
#include <vector>
/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n).
namespace noya {
/// @brief Generic rerooting DP on a tree.
/// merge combines neighbor contributions, ae(value, from, to) moves a
/// rooted value across an edge, and av(value, vertex) closes a vertex.
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);
int ds = 0;
for (const auto &adj : tr) {
ds += int(adj.size());
}
assert(ds == 2 * (n - 1));
std::vector<int> fa(n, -1);
std::vector<int> ord = {rt};
fa[rt] = rt;
for (int i = 0; i < int(ord.size()); i++) {
int u = ord[i];
for (int v : tr[u]) {
assert(0 <= v && v < n);
if (fa[v] == -1) {
fa[v] = u;
ord.push_back(v);
}
}
}
assert(int(ord.size()) == n);
std::vector<T> dn(n, uni);
for (int i = n - 1; i >= 0; i--) {
int u = ord[i];
T val = uni;
for (int v : tr[u]) {
if (fa[v] == u) {
val = mer(val, ae(dn[v], v, u));
}
}
dn[u] = av(val, u);
}
std::vector<T> up(n, uni);
std::vector<T> ans(n, uni);
for (int u : ord) {
int deg = int(tr[u].size());
std::vector<T> par(deg, uni);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
par[i] = v == fa[u] ? up[u] : ae(dn[v], v, u);
}
std::vector<T> pre(deg + 1, uni);
std::vector<T> suf(deg + 1, uni);
for (int i = 0; i < deg; i++) {
pre[i + 1] = mer(pre[i], par[i]);
}
for (int i = deg - 1; i >= 0; i--) {
suf[i] = mer(par[i], suf[i + 1]);
}
ans[u] = av(pre[deg], u);
for (int i = 0; i < deg; i++) {
int v = tr[u][i];
if (fa[v] == u) {
T wc = mer(pre[i], suf[i + 1]);
up[v] = ae(av(wc, u), u, v);
}
}
}
return ans;
}
} // namespace noya