tree_path_affine_sum.hpp¶
维护树上点的仿射函数，并查询一条有向路径上的函数复合结果。
Complexity: Time: O(n). Space: O(n), including O(n) recursion stack on a path.
AC 记录：tree_path_composite_sum。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n).
/// Space: O(n), including O(n) recursion stack on a path.
#include <cassert>
#include <utility>
#include <vector>
namespace noya {
/// @brief Sum edge-affine path transforms from every possible root. For an
/// oriented edge x-y, the message sent by y is the sum of path values whose
/// destinations lie on y's side after deleting the edge. Applying bx+c to all
/// of them gives b times that message plus c times the side size. One postorder
/// pass builds child-side messages; one preorder reroot pass supplies the
/// complementary parent-side message.
template <class T>
std::vector<T> tree_path_affine_sums(const std::vector<T> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<T, T>> &aff) {
int n = int(val.size());
assert(ee.size() == aff.size());
assert(int(ee.size()) == std::max(0, n - 1));
if (n == 0) {
return {};
}
struct edge {
int to;
T mul;
T add;
};
std::vector<std::vector<edge>> g(n);
for (int i = 0; i < n - 1; i++) {
auto [a, b] = ee[i];
auto [mul, add] = aff[i];
g[a].push_back({b, mul, add});
g[b].push_back({a, mul, add});
}
std::vector<int> siz(n);
std::vector<T> dn(n);
auto dfs = [&](auto &&self, int u, int fa) -> void {
siz[u] = 1;
dn[u] = val[u];
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
self(self, x.to, u);
siz[u] += siz[x.to];
dn[u] += x.mul * dn[x.to] + x.add * T(siz[x.to]);
}
};
dfs(dfs, 0, -1);
std::vector<T> ans(n);
ans[0] = dn[0];
auto df1 = [&](auto &&self, int u, int fa) -> void {
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
T cc = x.mul * dn[x.to] + x.add * T(siz[x.to]);
T up = ans[u] - cc;
T pc = x.mul * up + x.add * T(n - siz[x.to]);
ans[x.to] = dn[x.to] + pc;
self(self, x.to, u);
}
};
df1(df1, 0, -1);
return ans;
}
} // namespace noya
#ifndef NOYA_TREE_PATH_AFFINE_SUM_HPP
#define NOYA_TREE_PATH_AFFINE_SUM_HPP 1
/// @complexity Time: O(n).
/// Space: O(n), including O(n) recursion stack on a path.
#include <cassert>
#include <utility>
#include <vector>
namespace noya {
/// @brief Sum edge-affine path transforms from every possible root. For an
/// oriented edge x-y, the message sent by y is the sum of path values whose
/// destinations lie on y's side after deleting the edge. Applying bx+c to all
/// of them gives b times that message plus c times the side size. One postorder
/// pass builds child-side messages; one preorder reroot pass supplies the
/// complementary parent-side message.
template <class T>
std::vector<T> tree_path_affine_sums(const std::vector<T> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<T, T>> &aff) {
int n = int(val.size());
assert(ee.size() == aff.size());
assert(int(ee.size()) == std::max(0, n - 1));
if (n == 0) {
return {};
}
struct edge {
int to;
T mul;
T add;
};
std::vector<std::vector<edge>> g(n);
for (int i = 0; i < n - 1; i++) {
auto [a, b] = ee[i];
auto [mul, add] = aff[i];
g[a].push_back({b, mul, add});
g[b].push_back({a, mul, add});
}
std::vector<int> siz(n);
std::vector<T> dn(n);
auto dfs = [&](auto &&self, int u, int fa) -> void {
siz[u] = 1;
dn[u] = val[u];
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
self(self, x.to, u);
siz[u] += siz[x.to];
dn[u] += x.mul * dn[x.to] + x.add * T(siz[x.to]);
}
};
dfs(dfs, 0, -1);
std::vector<T> ans(n);
ans[0] = dn[0];
auto df1 = [&](auto &&self, int u, int fa) -> void {
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
T cc = x.mul * dn[x.to] + x.add * T(siz[x.to]);
T up = ans[u] - cc;
T pc = x.mul * up + x.add * T(n - siz[x.to]);
ans[x.to] = dn[x.to] + pc;
self(self, x.to, u);
}
};
df1(df1, 0, -1);
return ans;
}
} // namespace noya
#endif // NOYA_TREE_PATH_AFFINE_SUM_HPP
#include <cassert>
#include <utility>
#include <vector>
/// @complexity Time: O(n).
/// Space: O(n), including O(n) recursion stack on a path.
namespace noya {
/// @brief Sum edge-affine path transforms from every possible root. For an
/// oriented edge x-y, the message sent by y is the sum of path values whose
/// destinations lie on y's side after deleting the edge. Applying bx+c to all
/// of them gives b times that message plus c times the side size. One postorder
/// pass builds child-side messages; one preorder reroot pass supplies the
/// complementary parent-side message.
template <class T>
std::vector<T> tree_path_affine_sums(const std::vector<T> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<T, T>> &aff) {
int n = int(val.size());
assert(ee.size() == aff.size());
assert(int(ee.size()) == std::max(0, n - 1));
if (n == 0) {
return {};
}
struct edge {
int to;
T mul;
T add;
};
std::vector<std::vector<edge>> g(n);
for (int i = 0; i < n - 1; i++) {
auto [a, b] = ee[i];
auto [mul, add] = aff[i];
g[a].push_back({b, mul, add});
g[b].push_back({a, mul, add});
}
std::vector<int> siz(n);
std::vector<T> dn(n);
auto dfs = [&](auto &&self, int u, int fa) -> void {
siz[u] = 1;
dn[u] = val[u];
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
self(self, x.to, u);
siz[u] += siz[x.to];
dn[u] += x.mul * dn[x.to] + x.add * T(siz[x.to]);
}
};
dfs(dfs, 0, -1);
std::vector<T> ans(n);
ans[0] = dn[0];
auto df1 = [&](auto &&self, int u, int fa) -> void {
for (const edge &x : g[u]) {
if (x.to == fa) {
continue;
}
T cc = x.mul * dn[x.to] + x.add * T(siz[x.to]);
T up = ans[u] - cc;
T pc = x.mul * up + x.add * T(n - siz[x.to]);
ans[x.to] = dn[x.to] + pc;
self(self, x.to, u);
}
};
df1(df1, 0, -1);
return ans;
}
} // namespace noya