tree_path_affine_sum.hpp¶
静态计算每个终点收到的所有来源点值沿路径应用边仿射函数后的总和。两个方向使用同一边函数，靠近终点的函数最后应用；两趟递归缓存子树贡献，不提供更新或单条路径查询，支持零乘数与矩阵环。
Complexity: Time: \(O(n)\) scalar operations. Space: \(O(n)\), including \(O(n)\) recursive depth.
AC 记录：tree_path_composite_sum。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n) scalar operations. Space: O(n), including O(n) recursive depth.
#include <cassert>
#include <limits>
#include <utility>
#include <vector>
namespace noya {
/// @brief For every root r, sum all vertex values transformed along v -> r.
/// Edge i uses aff[i].first*x + aff[i].second in both directions, not inverse
/// maps. The edge nearest r is applied last; r contributes val[r] unchanged.
/// Static computation: no updates or single-path queries. Empty input returns {}.
/// ee and aff must describe the same undirected tree on [0,val.size()); n fits
/// int. Inputs are unchanged. T supports default/copy construction, T(int),
/// +, -, *, += with ring laws; integer conversion represents multiples of the
/// multiplicative identity. Multiplication need not commute; no division or
/// equality is required. All arithmetic intermediates must fit T.
/// Cache each child-to-parent transformed sum during postorder, then reuse it
/// when supplying the complementary parent side. Recursive depth is O(n).
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) {
assert(val.size() <= std::size_t(std::numeric_limits<int>::max()));
int n = int(val.size());
assert(ee.size() == aff.size());
assert(ee.size() == (val.empty() ? 0 : val.size() - 1));
if (!n) return {};
struct edge { int to, id; };
std::vector<std::vector<edge>> g(n);
std::vector<int> size(n);
for (auto [u, v] : ee) {
assert(0 <= u && u < n && 0 <= v && v < n);
++size[u]; ++size[v];
}
for (int u = 0; u < n; ++u) g[u].reserve(size[u]);
for (int i = 0; i < n - 1; ++i) {
auto [u, v] = ee[i];
g[u].push_back({v, i});
g[v].push_back({u, i});
}
auto answer = val;
std::vector<T> contribution(n);
auto down = [&](auto &self, int u, int parent) -> void {
size[u] = 1;
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
self(self, e.to, u);
size[u] += size[e.to];
contribution[e.to] = mul * answer[e.to] + add * T(size[e.to]);
answer[u] += contribution[e.to];
}
};
down(down, 0, -1);
auto reroot = [&](auto &self, int u, int parent) -> void {
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
T outside = answer[u] - contribution[e.to];
answer[e.to] += mul * outside + add * T(n - size[e.to]);
self(self, e.to, u);
}
};
reroot(reroot, 0, -1);
return answer;
}
} // namespace noya
// NOYA_TREE_PATH_AFFINE_SUM_HPP
#ifndef NOYA_TREE_PATH_AFFINE_SUM_HPP
#define NOYA_TREE_PATH_AFFINE_SUM_HPP 1
/// @complexity Time: O(n) scalar operations. Space: O(n), including O(n) recursive depth.
#include <cassert>
#include <limits>
#include <utility>
#include <vector>
namespace noya {
/// @brief For every root r, sum all vertex values transformed along v -> r.
/// Edge i uses aff[i].first*x + aff[i].second in both directions, not inverse
/// maps. The edge nearest r is applied last; r contributes val[r] unchanged.
/// Static computation: no updates or single-path queries. Empty input returns {}.
/// ee and aff must describe the same undirected tree on [0,val.size()); n fits
/// int. Inputs are unchanged. T supports default/copy construction, T(int),
/// +, -, *, += with ring laws; integer conversion represents multiples of the
/// multiplicative identity. Multiplication need not commute; no division or
/// equality is required. All arithmetic intermediates must fit T.
/// Cache each child-to-parent transformed sum during postorder, then reuse it
/// when supplying the complementary parent side. Recursive depth is O(n).
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) {
assert(val.size() <= std::size_t(std::numeric_limits<int>::max()));
int n = int(val.size());
assert(ee.size() == aff.size());
assert(ee.size() == (val.empty() ? 0 : val.size() - 1));
if (!n) return {};
struct edge { int to, id; };
std::vector<std::vector<edge>> g(n);
std::vector<int> size(n);
for (auto [u, v] : ee) {
assert(0 <= u && u < n && 0 <= v && v < n);
++size[u]; ++size[v];
}
for (int u = 0; u < n; ++u) g[u].reserve(size[u]);
for (int i = 0; i < n - 1; ++i) {
auto [u, v] = ee[i];
g[u].push_back({v, i});
g[v].push_back({u, i});
}
auto answer = val;
std::vector<T> contribution(n);
auto down = [&](auto &self, int u, int parent) -> void {
size[u] = 1;
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
self(self, e.to, u);
size[u] += size[e.to];
contribution[e.to] = mul * answer[e.to] + add * T(size[e.to]);
answer[u] += contribution[e.to];
}
};
down(down, 0, -1);
auto reroot = [&](auto &self, int u, int parent) -> void {
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
T outside = answer[u] - contribution[e.to];
answer[e.to] += mul * outside + add * T(n - size[e.to]);
self(self, e.to, u);
}
};
reroot(reroot, 0, -1);
return answer;
}
} // namespace noya
#endif // NOYA_TREE_PATH_AFFINE_SUM_HPP
/// @complexity Time: O(n) scalar operations. Space: O(n), including O(n) recursive depth.
#include <cassert>
#include <limits>
#include <utility>
#include <vector>
namespace noya {
/// @brief For every root r, sum all vertex values transformed along v -> r.
/// Edge i uses aff[i].first*x + aff[i].second in both directions, not inverse
/// maps. The edge nearest r is applied last; r contributes val[r] unchanged.
/// Static computation: no updates or single-path queries. Empty input returns {}.
/// ee and aff must describe the same undirected tree on [0,val.size()); n fits
/// int. Inputs are unchanged. T supports default/copy construction, T(int),
/// +, -, *, += with ring laws; integer conversion represents multiples of the
/// multiplicative identity. Multiplication need not commute; no division or
/// equality is required. All arithmetic intermediates must fit T.
/// Cache each child-to-parent transformed sum during postorder, then reuse it
/// when supplying the complementary parent side. Recursive depth is O(n).
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) {
assert(val.size() <= std::size_t(std::numeric_limits<int>::max()));
int n = int(val.size());
assert(ee.size() == aff.size());
assert(ee.size() == (val.empty() ? 0 : val.size() - 1));
if (!n) return {};
struct edge { int to, id; };
std::vector<std::vector<edge>> g(n);
std::vector<int> size(n);
for (auto [u, v] : ee) {
assert(0 <= u && u < n && 0 <= v && v < n);
++size[u]; ++size[v];
}
for (int u = 0; u < n; ++u) g[u].reserve(size[u]);
for (int i = 0; i < n - 1; ++i) {
auto [u, v] = ee[i];
g[u].push_back({v, i});
g[v].push_back({u, i});
}
auto answer = val;
std::vector<T> contribution(n);
auto down = [&](auto &self, int u, int parent) -> void {
size[u] = 1;
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
self(self, e.to, u);
size[u] += size[e.to];
contribution[e.to] = mul * answer[e.to] + add * T(size[e.to]);
answer[u] += contribution[e.to];
}
};
down(down, 0, -1);
auto reroot = [&](auto &self, int u, int parent) -> void {
for (const auto &e : g[u]) if (e.to != parent) {
const auto &[mul, add] = aff[e.id];
T outside = answer[u] - contribution[e.to];
answer[e.to] += mul * outside + add * T(n - size[e.to]);
self(self, e.to, u);
}
};
reroot(reroot, 0, -1);
return answer;
}
} // namespace noya
// NOYA_TREE_PATH_AFFINE_SUM_HPP