tree_path_affine_sum.hpp¶
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.
Verified by tree_path_composite_sum.
维护树上点的仿射函数，并查询一条有向路径上的函数复合结果。
Implementation¶
#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> &vertex_values,
const std::vector<std::pair<int, int>> &edge_endpoints,
const std::vector<std::pair<T, T>> &edge_affine) {
int n = int(vertex_values.size());
assert(edge_endpoints.size() == edge_affine.size());
assert(int(edge_endpoints.size()) == std::max(0, n - 1));
if (n == 0) {
return {};
}
struct edge {
int to;
T multiplier;
T addition;
};
std::vector<std::vector<edge>> graph(n);
for (int i = 0; i < n - 1; i++) {
auto [first, second] = edge_endpoints[i];
auto [multiplier, addition] = edge_affine[i];
graph[first].push_back({second, multiplier, addition});
graph[second].push_back({first, multiplier, addition});
}
std::vector<int> subtree_size(n);
std::vector<T> downward(n);
auto build = [&](auto &&self, int vertex, int parent) -> void {
subtree_size[vertex] = 1;
downward[vertex] = vertex_values[vertex];
for (const edge &item : graph[vertex]) {
if (item.to == parent) {
continue;
}
self(self, item.to, vertex);
subtree_size[vertex] += subtree_size[item.to];
downward[vertex] += item.multiplier * downward[item.to] +
item.addition * T(subtree_size[item.to]);
}
};
build(build, 0, -1);
std::vector<T> answer(n);
answer[0] = downward[0];
auto reroot = [&](auto &&self, int vertex, int parent) -> void {
for (const edge &item : graph[vertex]) {
if (item.to == parent) {
continue;
}
T child_contribution =
item.multiplier * downward[item.to] +
item.addition * T(subtree_size[item.to]);
T parent_side_message = answer[vertex] - child_contribution;
T parent_contribution =
item.multiplier * parent_side_message +
item.addition * T(n - subtree_size[item.to]);
answer[item.to] = downward[item.to] + parent_contribution;
self(self, item.to, vertex);
}
};
reroot(reroot, 0, -1);
return answer;
}
} // 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> &vertex_values,
const std::vector<std::pair<int, int>> &edge_endpoints,
const std::vector<std::pair<T, T>> &edge_affine) {
int n = int(vertex_values.size());
assert(edge_endpoints.size() == edge_affine.size());
assert(int(edge_endpoints.size()) == std::max(0, n - 1));
if (n == 0) {
return {};
}
struct edge {
int to;
T multiplier;
T addition;
};
std::vector<std::vector<edge>> graph(n);
for (int i = 0; i < n - 1; i++) {
auto [first, second] = edge_endpoints[i];
auto [multiplier, addition] = edge_affine[i];
graph[first].push_back({second, multiplier, addition});
graph[second].push_back({first, multiplier, addition});
}
std::vector<int> subtree_size(n);
std::vector<T> downward(n);
auto build = [&](auto &&self, int vertex, int parent) -> void {
subtree_size[vertex] = 1;
downward[vertex] = vertex_values[vertex];
for (const edge &item : graph[vertex]) {
if (item.to == parent) {
continue;
}
self(self, item.to, vertex);
subtree_size[vertex] += subtree_size[item.to];
downward[vertex] += item.multiplier * downward[item.to] +
item.addition * T(subtree_size[item.to]);
}
};
build(build, 0, -1);
std::vector<T> answer(n);
answer[0] = downward[0];
auto reroot = [&](auto &&self, int vertex, int parent) -> void {
for (const edge &item : graph[vertex]) {
if (item.to == parent) {
continue;
}
T child_contribution =
item.multiplier * downward[item.to] +
item.addition * T(subtree_size[item.to]);
T parent_side_message = answer[vertex] - child_contribution;
T parent_contribution =
item.multiplier * parent_side_message +
item.addition * T(n - subtree_size[item.to]);
answer[item.to] = downward[item.to] + parent_contribution;
self(self, item.to, vertex);
}
};
reroot(reroot, 0, -1);
return answer;
}
} // namespace noya