dynamic_tree_path_affine_sum.hpp¶
Maintain sums of edge-affine transforms from any chosen root. Every original edge is split by an edge-node carrying its affine map. A top tree cluster stores the affine transform along its boundary path, the sum of all values entering that path, and their count. Compress composes paths; rake adds independent branches. Preferred-path expose therefore propagates a point update to the complete rerooted aggregate in logarithmic time.
Verified by point_set_tree_path_composite_sum, point_set_tree_path_composite_sum_fixed_root.
在动态森林中连边断边、修改点仿射函数，并查询有向路径的函数复合。
Implementation¶
#ifndef NOYA_DYNAMIC_TREE_PATH_AFFINE_SUM_HPP
#define NOYA_DYNAMIC_TREE_PATH_AFFINE_SUM_HPP 1
/// @complexity Time: O(n log n) construction and O(log n) amortized per
/// vertex update, edge update, or root query. Space: O(n).
#include "noya/dynamic_rerooting_top_tree.hpp"
#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>
namespace noya {
namespace dynamic_tree_path_affine_sum_internal {
template <class Mint> struct traits {
struct path {
Mint multiplier{};
Mint addition{};
Mint sum{};
Mint count{};
};
struct point {
Mint sum{};
Mint count{};
};
struct info {
bool is_vertex = true;
Mint first{};
Mint second{};
};
static path vertex(const info &item) {
if (item.is_vertex) {
return {Mint(1), Mint(0), item.first, Mint(1)};
}
return {item.first, item.second, Mint(0), Mint(0)};
}
static path compress(const path &parent, const path &child) {
return {parent.multiplier * child.multiplier,
parent.multiplier * child.addition + parent.addition,
parent.sum + parent.multiplier * child.sum +
parent.addition * child.count,
parent.count + child.count};
}
static point rake(const point &first, const point &second) {
return {first.sum + second.sum, first.count + second.count};
}
static point add_edge(const path &cluster) {
return {cluster.sum, cluster.count};
}
static path add_vertex(const point &children, const info &item) {
if (item.is_vertex) {
return {Mint(1), Mint(0), children.sum + item.first,
children.count + Mint(1)};
}
return {item.first, item.second,
children.sum * item.first + children.count * item.second,
children.count};
}
};
} // namespace dynamic_tree_path_affine_sum_internal
/// @brief Maintain sums of edge-affine transforms from any chosen root.
/// Every original edge is split by an edge-node carrying its affine map. A top
/// tree cluster stores the affine transform along its boundary path, the sum
/// of all values entering that path, and their count. Compress composes paths;
/// rake adds independent branches. Preferred-path expose therefore propagates
/// a point update to the complete rerooted aggregate in logarithmic time.
template <class Mint> class dynamic_tree_path_affine_sum {
using trait = dynamic_tree_path_affine_sum_internal::traits<Mint>;
using path = typename trait::path;
using point = typename trait::point;
using info = typename trait::info;
using dp_type =
dynamic_rerooting_dp<path, point, info, trait::vertex, trait::compress,
trait::rake, trait::add_edge, trait::add_vertex>;
int vertex_count;
dp_type dp;
static std::vector<info>
make_info(const std::vector<Mint> &vertex_values,
const std::vector<std::pair<Mint, Mint>> &edge_affine) {
std::vector<info> result(vertex_values.size() + edge_affine.size());
for (int vertex = 0; vertex < int(vertex_values.size()); vertex++) {
result[vertex] = {true, vertex_values[vertex], Mint(0)};
}
for (int edge = 0; edge < int(edge_affine.size()); edge++) {
result[vertex_values.size() + edge] =
{false, edge_affine[edge].first, edge_affine[edge].second};
}
return result;
}
public:
dynamic_tree_path_affine_sum(
const std::vector<Mint> &vertex_values,
const std::vector<std::pair<int, int>> &edge_endpoints,
const std::vector<std::pair<Mint, Mint>> &edge_affine)
: vertex_count(int(vertex_values.size())),
dp(make_info(vertex_values, edge_affine)) {
assert(vertex_count >= 1);
assert(int(edge_endpoints.size()) == vertex_count - 1);
assert(edge_affine.size() == edge_endpoints.size());
for (int edge = 0; edge < vertex_count - 1; edge++) {
auto [first, second] = edge_endpoints[edge];
assert(0 <= first && first < vertex_count && 0 <= second &&
second < vertex_count);
dp.link(vertex_count + edge, first);
dp.link(vertex_count + edge, second);
}
}
void set_vertex(int vertex, Mint value) {
assert(0 <= vertex && vertex < vertex_count);
dp.set_info(vertex, {true, value, Mint(0)});
}
void set_edge(int edge, Mint multiplier, Mint addition) {
assert(0 <= edge && edge < vertex_count - 1);
dp.set_info(vertex_count + edge,
{false, multiplier, addition});
}
Mint query(int root) {
assert(0 <= root && root < vertex_count);
return dp.query(root).sum;
}
};
} // namespace noya
#endif // NOYA_DYNAMIC_TREE_PATH_AFFINE_SUM_HPP
#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>
/// @complexity Time: O(n log n) construction and O(log n) amortized per
/// vertex update, edge update, or root query. Space: O(n).
/// @complexity Time: O(log n) amortized per link, cut, point update, reroot,
/// or aggregate query. Space: O(n).
namespace noya {
namespace dynamic_rerooting_internal {
template <class Point, auto rake> class dashed_splay_tree {
public:
struct node {
node *left = nullptr;
node *right = nullptr;
node *parent = nullptr;
Point key;
Point sum;
explicit node(const Point &value) : key(value), sum(value) {}
};
using node_pointer = node *;
private:
static void update(node_pointer current) {
current->sum = current->key;
if (current->left) {
current->sum = rake(current->sum, current->left->sum);
}
if (current->right) {
current->sum = rake(current->sum, current->right->sum);
}
}
static void rotate_right(node_pointer current) {
node_pointer parent = current->parent;
node_pointer grandparent = parent->parent;
parent->left = current->right;
if (current->right) {
current->right->parent = parent;
}
current->right = parent;
parent->parent = current;
update(parent);
update(current);
current->parent = grandparent;
if (grandparent) {
if (grandparent->left == parent) {
grandparent->left = current;
}
if (grandparent->right == parent) {
grandparent->right = current;
}
}
}
static void rotate_left(node_pointer current) {
node_pointer parent = current->parent;
node_pointer grandparent = parent->parent;
parent->right = current->left;
if (current->left) {
current->left->parent = parent;
}
current->left = parent;
parent->parent = current;
update(parent);
update(current);
current->parent = grandparent;
if (grandparent) {
if (grandparent->left == parent) {
grandparent->left = current;
}
if (grandparent->right == parent) {
grandparent->right = current;
}
}
}
static node_pointer rightmost(node_pointer current) {
while (current->right) {
current = current->right;
}
return current;
}
public:
static void splay(node_pointer current) {
while (current->parent) {
node_pointer parent = current->parent;
if (!parent->parent) {
if (parent->left == current) {
rotate_right(current);
} else {
rotate_left(current);
}
} else {
node_pointer grandparent = parent->parent;
if (grandparent->left == parent) {
if (parent->left == current) {
rotate_right(parent);
rotate_right(current);
} else {
rotate_left(current);
rotate_right(current);
}
} else if (parent->right == current) {
rotate_left(parent);
rotate_left(current);
} else {
rotate_right(current);
rotate_left(current);
}
}
}
}
static node_pointer insert(node_pointer root, const Point &value) {
if (!root) {
return new node(value);
}
node_pointer last = rightmost(root);
splay(last);
node_pointer inserted = new node(value);
inserted->parent = last;
last->right = inserted;
update(last);
splay(inserted);
return inserted;
}
static node_pointer erase(node_pointer current) {
splay(current);
node_pointer left = current->left;
node_pointer right = current->right;
delete current;
if (!left) {
if (right) {
right->parent = nullptr;
}
return right;
}
if (!right) {
left->parent = nullptr;
return left;
}
left->parent = nullptr;
node_pointer root = rightmost(left);
splay(root);
root->right = right;
right->parent = root;
update(root);
return root;
}
};
} // namespace dynamic_rerooting_internal
/// @brief Dynamic rerooting DP on a top tree represented by preferred paths.
/// `compress` joins consecutive path clusters, while commutative `rake`
/// combines light subtrees. `add_edge` changes a path cluster into a point
/// cluster and `add_vertex` attaches all light clusters to one path vertex.
/// Expose moves preferred edges between the path splay and the dashed-edge
/// splay, so both forward and reversed path aggregates stay current.
template <class Path, class Point, class Info, auto vertex, auto compress,
auto rake, auto add_edge, auto add_vertex>
class dynamic_rerooting_top_tree {
using dashed_tree =
dynamic_rerooting_internal::dashed_splay_tree<Point, rake>;
struct node {
node *left = nullptr;
node *right = nullptr;
node *parent = nullptr;
Info info;
Path key{};
Path forward{};
Path backward{};
typename dashed_tree::node_pointer light = nullptr;
typename dashed_tree::node_pointer light_position = nullptr;
bool reversed = false;
explicit node(const Info &value) : info(value) {}
bool is_auxiliary_root() const {
return !parent || (parent->left != this && parent->right != this);
}
};
public:
using node_pointer = node *;
private:
static void toggle(node_pointer current) {
std::swap(current->left, current->right);
std::swap(current->forward, current->backward);
current->reversed ^= true;
}
static void push(node_pointer current) {
if (!current->reversed) {
return;
}
if (current->left) {
toggle(current->left);
}
if (current->right) {
toggle(current->right);
}
current->reversed = false;
}
static void update(node_pointer current) {
Path key = current->light
? add_vertex(current->light->sum, current->info)
: vertex(current->info);
Path forward = key;
Path backward = key;
if (current->left) {
forward = compress(current->left->forward, forward);
backward = compress(backward, current->left->backward);
}
if (current->right) {
forward = compress(forward, current->right->forward);
backward = compress(current->right->backward, backward);
}
current->key = key;
current->forward = forward;
current->backward = backward;
}
static void rotate_right(node_pointer current) {
node_pointer parent = current->parent;
node_pointer grandparent = parent->parent;
push(parent);
push(current);
parent->left = current->right;
if (current->right) {
current->right->parent = parent;
}
current->right = parent;
parent->parent = current;
update(parent);
update(current);
current->parent = grandparent;
if (grandparent) {
if (grandparent->left == parent) {
grandparent->left = current;
}
if (grandparent->right == parent) {
grandparent->right = current;
}
}
}
static void rotate_left(node_pointer current) {
node_pointer parent = current->parent;
node_pointer grandparent = parent->parent;
push(parent);
push(current);
parent->right = current->left;
if (current->left) {
current->left->parent = parent;
}
current->left = parent;
parent->parent = current;
update(parent);
update(current);
current->parent = grandparent;
if (grandparent) {
if (grandparent->left == parent) {
grandparent->left = current;
}
if (grandparent->right == parent) {
grandparent->right = current;
}
}
}
static void splay(node_pointer current) {
push(current);
{
node_pointer root = current;
while (!root->is_auxiliary_root()) {
root = root->parent;
}
current->light_position = root->light_position;
if (current != root) {
root->light_position = nullptr;
}
}
while (!current->is_auxiliary_root()) {
node_pointer parent = current->parent;
if (parent->is_auxiliary_root()) {
push(parent);
push(current);
if (parent->left == current) {
rotate_right(current);
} else {
rotate_left(current);
}
} else {
node_pointer grandparent = parent->parent;
push(grandparent);
push(parent);
push(current);
if (grandparent->left == parent) {
if (parent->left == current) {
rotate_right(parent);
rotate_right(current);
} else {
rotate_left(current);
rotate_right(current);
}
} else if (parent->right == current) {
rotate_left(parent);
rotate_left(current);
} else {
rotate_right(current);
rotate_left(current);
}
}
}
}
static node_pointer expose(node_pointer current) {
node_pointer preferred = nullptr;
for (node_pointer ancestor = current; ancestor; ancestor = ancestor->parent) {
splay(ancestor);
if (ancestor->right) {
ancestor->light =
dashed_tree::insert(ancestor->light,
add_edge(ancestor->right->forward));
ancestor->right->light_position = ancestor->light;
}
ancestor->right = preferred;
if (preferred) {
dashed_tree::splay(preferred->light_position);
push(preferred);
ancestor->light = dashed_tree::erase(preferred->light_position);
}
update(ancestor);
preferred = ancestor;
}
splay(current);
return preferred;
}
public:
node_pointer make_node(const Info &info) {
node_pointer result = new node(info);
update(result);
return result;
}
void evert(node_pointer current) {
expose(current);
toggle(current);
push(current);
}
void link(node_pointer child, node_pointer parent) {
expose(parent);
expose(child);
child->parent = parent;
parent->right = child;
update(parent);
}
void cut(node_pointer child) {
expose(child);
node_pointer parent = child->left;
assert(parent);
child->left = nullptr;
parent->parent = nullptr;
update(child);
}
void set_info(node_pointer current, const Info &info) {
expose(current);
current->info = info;
update(current);
}
Info get_info(node_pointer current) const { return current->info; }
Path query(node_pointer root) {
evert(root);
return root->forward;
}
Path query_subtree(node_pointer root, node_pointer subtree_root) {
evert(root);
expose(subtree_root);
node_pointer prefix = subtree_root->left;
subtree_root->left = nullptr;
update(subtree_root);
Path result = subtree_root->forward;
subtree_root->left = prefix;
update(subtree_root);
return result;
}
};
/// @brief Index-based wrapper around `dynamic_rerooting_top_tree`.
template <class Path, class Point, class Info, auto vertex, auto compress,
auto rake, auto add_edge, auto add_vertex>
class dynamic_rerooting_dp {
using top_tree_type = dynamic_rerooting_top_tree<
Path, Point, Info, vertex, compress, rake, add_edge, add_vertex>;
top_tree_type tree;
std::vector<typename top_tree_type::node_pointer> nodes;
public:
explicit dynamic_rerooting_dp(const std::vector<Info> &info)
: nodes(info.size()) {
for (int i = 0; i < int(info.size()); i++) {
nodes[i] = tree.make_node(info[i]);
}
}
void link(int first, int second) {
tree.evert(nodes[first]);
tree.link(nodes[first], nodes[second]);
}
void cut(int first, int second) {
tree.evert(nodes[first]);
tree.cut(nodes[second]);
}
void set_info(int index, const Info &info) {
tree.set_info(nodes[index], info);
}
Info get_info(int index) const { return tree.get_info(nodes[index]); }
Path query(int root) { return tree.query(nodes[root]); }
Path query_subtree(int root, int subtree_root) {
return tree.query_subtree(nodes[root], nodes[subtree_root]);
}
};
} // namespace noya
namespace noya {
namespace dynamic_tree_path_affine_sum_internal {
template <class Mint> struct traits {
struct path {
Mint multiplier{};
Mint addition{};
Mint sum{};
Mint count{};
};
struct point {
Mint sum{};
Mint count{};
};
struct info {
bool is_vertex = true;
Mint first{};
Mint second{};
};
static path vertex(const info &item) {
if (item.is_vertex) {
return {Mint(1), Mint(0), item.first, Mint(1)};
}
return {item.first, item.second, Mint(0), Mint(0)};
}
static path compress(const path &parent, const path &child) {
return {parent.multiplier * child.multiplier,
parent.multiplier * child.addition + parent.addition,
parent.sum + parent.multiplier * child.sum +
parent.addition * child.count,
parent.count + child.count};
}
static point rake(const point &first, const point &second) {
return {first.sum + second.sum, first.count + second.count};
}
static point add_edge(const path &cluster) {
return {cluster.sum, cluster.count};
}
static path add_vertex(const point &children, const info &item) {
if (item.is_vertex) {
return {Mint(1), Mint(0), children.sum + item.first,
children.count + Mint(1)};
}
return {item.first, item.second,
children.sum * item.first + children.count * item.second,
children.count};
}
};
} // namespace dynamic_tree_path_affine_sum_internal
/// @brief Maintain sums of edge-affine transforms from any chosen root.
/// Every original edge is split by an edge-node carrying its affine map. A top
/// tree cluster stores the affine transform along its boundary path, the sum
/// of all values entering that path, and their count. Compress composes paths;
/// rake adds independent branches. Preferred-path expose therefore propagates
/// a point update to the complete rerooted aggregate in logarithmic time.
template <class Mint> class dynamic_tree_path_affine_sum {
using trait = dynamic_tree_path_affine_sum_internal::traits<Mint>;
using path = typename trait::path;
using point = typename trait::point;
using info = typename trait::info;
using dp_type =
dynamic_rerooting_dp<path, point, info, trait::vertex, trait::compress,
trait::rake, trait::add_edge, trait::add_vertex>;
int vertex_count;
dp_type dp;
static std::vector<info>
make_info(const std::vector<Mint> &vertex_values,
const std::vector<std::pair<Mint, Mint>> &edge_affine) {
std::vector<info> result(vertex_values.size() + edge_affine.size());
for (int vertex = 0; vertex < int(vertex_values.size()); vertex++) {
result[vertex] = {true, vertex_values[vertex], Mint(0)};
}
for (int edge = 0; edge < int(edge_affine.size()); edge++) {
result[vertex_values.size() + edge] =
{false, edge_affine[edge].first, edge_affine[edge].second};
}
return result;
}
public:
dynamic_tree_path_affine_sum(
const std::vector<Mint> &vertex_values,
const std::vector<std::pair<int, int>> &edge_endpoints,
const std::vector<std::pair<Mint, Mint>> &edge_affine)
: vertex_count(int(vertex_values.size())),
dp(make_info(vertex_values, edge_affine)) {
assert(vertex_count >= 1);
assert(int(edge_endpoints.size()) == vertex_count - 1);
assert(edge_affine.size() == edge_endpoints.size());
for (int edge = 0; edge < vertex_count - 1; edge++) {
auto [first, second] = edge_endpoints[edge];
assert(0 <= first && first < vertex_count && 0 <= second &&
second < vertex_count);
dp.link(vertex_count + edge, first);
dp.link(vertex_count + edge, second);
}
}
void set_vertex(int vertex, Mint value) {
assert(0 <= vertex && vertex < vertex_count);
dp.set_info(vertex, {true, value, Mint(0)});
}
void set_edge(int edge, Mint multiplier, Mint addition) {
assert(0 <= edge && edge < vertex_count - 1);
dp.set_info(vertex_count + edge,
{false, multiplier, addition});
}
Mint query(int root) {
assert(0 <= root && root < vertex_count);
return dp.query(root).sum;
}
};
} // namespace noya