dynamic_tree_path_affine_sum.hpp¶
在动态森林中连边断边、修改点仿射函数，并查询有向路径的函数复合。
Complexity: Time: O(n log n) construction and O(log n) amortized per vertex update, edge update, or root query. Space: O(n).
AC 记录：point_set_tree_path_composite_sum, point_set_tree_path_composite_sum_fixed_root。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @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 mul{};
Mint add{};
Mint sum{};
Mint cnt{};
};
struct point {
Mint sum{};
Mint cnt{};
};
struct info {
bool iv = true;
Mint a{};
Mint b{};
};
static path vertex(const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), x.a, Mint(1)};
}
return {x.a, x.b, Mint(0), Mint(0)};
}
static path compress(const path &fa, const path &v) {
return {fa.mul * v.mul, fa.mul * v.add + fa.add,
fa.sum + fa.mul * v.sum + fa.add * v.cnt, fa.cnt + v.cnt};
}
static point rake(const point &a, const point &b) {
return {a.sum + b.sum, a.cnt + b.cnt};
}
static point add_edge(const path &cl) { return {cl.sum, cl.cnt}; }
static path add_vertex(const point &ch, const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), ch.sum + x.a, ch.cnt + Mint(1)};
}
return {x.a, x.b, ch.sum * x.a + ch.cnt * x.b, ch.cnt};
}
};
} // 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 n;
dp_type dp;
static std::vector<info>
make_info(const std::vector<Mint> &val,
const std::vector<std::pair<Mint, Mint>> &aff) {
std::vector<info> res(val.size() + aff.size());
for (int u = 0; u < int(val.size()); u++) {
res[u] = {true, val[u], Mint(0)};
}
for (int e = 0; e < int(aff.size()); e++) {
res[val.size() + e] = {false, aff[e].first, aff[e].second};
}
return res;
}
public:
dynamic_tree_path_affine_sum(const std::vector<Mint> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<Mint, Mint>> &aff)
: n(int(val.size())), dp(make_info(val, aff)) {
assert(n >= 1);
assert(int(ee.size()) == n - 1);
assert(aff.size() == ee.size());
for (int e = 0; e < n - 1; e++) {
auto [a, b] = ee[e];
assert(0 <= a && a < n && 0 <= b && b < n);
dp.link(n + e, a);
dp.link(n + e, b);
}
}
void set_vertex(int u, Mint va1) {
assert(0 <= u && u < n);
dp.set_info(u, {true, va1, Mint(0)});
}
void set_edge(int e, Mint mul, Mint add) {
assert(0 <= e && e < n - 1);
dp.set_info(n + e, {false, mul, add});
}
Mint query(int rt) {
assert(0 <= rt && rt < n);
return dp.query(rt).sum;
}
};
} // namespace noya
#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 mul{};
Mint add{};
Mint sum{};
Mint cnt{};
};
struct point {
Mint sum{};
Mint cnt{};
};
struct info {
bool iv = true;
Mint a{};
Mint b{};
};
static path vertex(const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), x.a, Mint(1)};
}
return {x.a, x.b, Mint(0), Mint(0)};
}
static path compress(const path &fa, const path &v) {
return {fa.mul * v.mul, fa.mul * v.add + fa.add,
fa.sum + fa.mul * v.sum + fa.add * v.cnt, fa.cnt + v.cnt};
}
static point rake(const point &a, const point &b) {
return {a.sum + b.sum, a.cnt + b.cnt};
}
static point add_edge(const path &cl) { return {cl.sum, cl.cnt}; }
static path add_vertex(const point &ch, const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), ch.sum + x.a, ch.cnt + Mint(1)};
}
return {x.a, x.b, ch.sum * x.a + ch.cnt * x.b, ch.cnt};
}
};
} // 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 n;
dp_type dp;
static std::vector<info>
make_info(const std::vector<Mint> &val,
const std::vector<std::pair<Mint, Mint>> &aff) {
std::vector<info> res(val.size() + aff.size());
for (int u = 0; u < int(val.size()); u++) {
res[u] = {true, val[u], Mint(0)};
}
for (int e = 0; e < int(aff.size()); e++) {
res[val.size() + e] = {false, aff[e].first, aff[e].second};
}
return res;
}
public:
dynamic_tree_path_affine_sum(const std::vector<Mint> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<Mint, Mint>> &aff)
: n(int(val.size())), dp(make_info(val, aff)) {
assert(n >= 1);
assert(int(ee.size()) == n - 1);
assert(aff.size() == ee.size());
for (int e = 0; e < n - 1; e++) {
auto [a, b] = ee[e];
assert(0 <= a && a < n && 0 <= b && b < n);
dp.link(n + e, a);
dp.link(n + e, b);
}
}
void set_vertex(int u, Mint va1) {
assert(0 <= u && u < n);
dp.set_info(u, {true, va1, Mint(0)});
}
void set_edge(int e, Mint mul, Mint add) {
assert(0 <= e && e < n - 1);
dp.set_info(n + e, {false, mul, add});
}
Mint query(int rt) {
assert(0 <= rt && rt < n);
return dp.query(rt).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 rk> class dashed_splay_tree {
public:
struct node {
node *l = nullptr;
node *r = nullptr;
node *fa = nullptr;
Point key;
Point sum;
explicit node(const Point &val) : key(val), sum(val) {}
};
using node_pointer = node *;
private:
static void update(node_pointer cur) {
cur->sum = cur->key;
if (cur->l) {
cur->sum = rk(cur->sum, cur->l->sum);
}
if (cur->r) {
cur->sum = rk(cur->sum, cur->r->sum);
}
}
static void rotate_right(node_pointer cur) {
node_pointer fa = cur->fa;
node_pointer gf = fa->fa;
fa->l = cur->r;
if (cur->r) {
cur->r->fa = fa;
}
cur->r = fa;
fa->fa = cur;
update(fa);
update(cur);
cur->fa = gf;
if (gf) {
if (gf->l == fa) {
gf->l = cur;
}
if (gf->r == fa) {
gf->r = cur;
}
}
}
static void rotate_left(node_pointer cur) {
node_pointer fa = cur->fa;
node_pointer gf = fa->fa;
fa->r = cur->l;
if (cur->l) {
cur->l->fa = fa;
}
cur->l = fa;
fa->fa = cur;
update(fa);
update(cur);
cur->fa = gf;
if (gf) {
if (gf->l == fa) {
gf->l = cur;
}
if (gf->r == fa) {
gf->r = cur;
}
}
}
static node_pointer rightmost(node_pointer cur) {
while (cur->r) {
cur = cur->r;
}
return cur;
}
public:
static void splay(node_pointer cur) {
while (cur->fa) {
node_pointer fa = cur->fa;
if (!fa->fa) {
if (fa->l == cur) {
rotate_right(cur);
} else {
rotate_left(cur);
}
} else {
node_pointer gf = fa->fa;
if (gf->l == fa) {
if (fa->l == cur) {
rotate_right(fa);
rotate_right(cur);
} else {
rotate_left(cur);
rotate_right(cur);
}
} else if (fa->r == cur) {
rotate_left(fa);
rotate_left(cur);
} else {
rotate_right(cur);
rotate_left(cur);
}
}
}
}
static node_pointer insert(node_pointer rt, const Point &val) {
if (!rt) {
return new node(val);
}
node_pointer lst = rightmost(rt);
splay(lst);
node_pointer ins = new node(val);
ins->fa = lst;
lst->r = ins;
update(lst);
splay(ins);
return ins;
}
static node_pointer erase(node_pointer cur) {
splay(cur);
node_pointer l = cur->l;
node_pointer r = cur->r;
delete cur;
if (!l) {
if (r) {
r->fa = nullptr;
}
return r;
}
if (!r) {
l->fa = nullptr;
return l;
}
l->fa = nullptr;
node_pointer rt = rightmost(l);
splay(rt);
rt->r = r;
r->fa = rt;
update(rt);
return rt;
}
};
} // namespace dynamic_rerooting_internal
/// @brief Dynamic rerooting DP on a top tree represented by preferred paths.
/// `cmp` joins consecutive path clusters, while commutative `rk`
/// combines light subtrees. `ae` changes a path cluster into a point
/// cluster and `av` 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 vtx, auto cmp, auto rk,
auto ae, auto av>
class dynamic_rerooting_top_tree {
using dashed_tree = dynamic_rerooting_internal::dashed_splay_tree<Point, rk>;
struct node {
node *l = nullptr;
node *r = nullptr;
node *fa = nullptr;
Info inf;
Path key{};
Path fwd{};
Path bwd{};
typename dashed_tree::node_pointer lt = nullptr;
typename dashed_tree::node_pointer lp = nullptr;
bool rev = false;
explicit node(const Info &val) : inf(val) {}
bool is_auxiliary_root() const {
return !fa || (fa->l != this && fa->r != this);
}
};
public:
using node_pointer = node *;
private:
static void toggle(node_pointer cur) {
std::swap(cur->l, cur->r);
std::swap(cur->fwd, cur->bwd);
cur->rev ^= true;
}
static void push(node_pointer cur) {
if (!cur->rev) {
return;
}
if (cur->l) {
toggle(cur->l);
}
if (cur->r) {
toggle(cur->r);
}
cur->rev = false;
}
static void update(node_pointer cur) {
Path key = cur->lt ? av(cur->lt->sum, cur->inf) : vtx(cur->inf);
Path fwd = key;
Path bwd = key;
if (cur->l) {
fwd = cmp(cur->l->fwd, fwd);
bwd = cmp(bwd, cur->l->bwd);
}
if (cur->r) {
fwd = cmp(fwd, cur->r->fwd);
bwd = cmp(cur->r->bwd, bwd);
}
cur->key = key;
cur->fwd = fwd;
cur->bwd = bwd;
}
static void rotate_right(node_pointer cur) {
node_pointer fa = cur->fa;
node_pointer gf = fa->fa;
push(fa);
push(cur);
fa->l = cur->r;
if (cur->r) {
cur->r->fa = fa;
}
cur->r = fa;
fa->fa = cur;
update(fa);
update(cur);
cur->fa = gf;
if (gf) {
if (gf->l == fa) {
gf->l = cur;
}
if (gf->r == fa) {
gf->r = cur;
}
}
}
static void rotate_left(node_pointer cur) {
node_pointer fa = cur->fa;
node_pointer gf = fa->fa;
push(fa);
push(cur);
fa->r = cur->l;
if (cur->l) {
cur->l->fa = fa;
}
cur->l = fa;
fa->fa = cur;
update(fa);
update(cur);
cur->fa = gf;
if (gf) {
if (gf->l == fa) {
gf->l = cur;
}
if (gf->r == fa) {
gf->r = cur;
}
}
}
static void splay(node_pointer cur) {
push(cur);
{
node_pointer rt = cur;
while (!rt->is_auxiliary_root()) {
rt = rt->fa;
}
cur->lp = rt->lp;
if (cur != rt) {
rt->lp = nullptr;
}
}
while (!cur->is_auxiliary_root()) {
node_pointer fa = cur->fa;
if (fa->is_auxiliary_root()) {
push(fa);
push(cur);
if (fa->l == cur) {
rotate_right(cur);
} else {
rotate_left(cur);
}
} else {
node_pointer gf = fa->fa;
push(gf);
push(fa);
push(cur);
if (gf->l == fa) {
if (fa->l == cur) {
rotate_right(fa);
rotate_right(cur);
} else {
rotate_left(cur);
rotate_right(cur);
}
} else if (fa->r == cur) {
rotate_left(fa);
rotate_left(cur);
} else {
rotate_right(cur);
rotate_left(cur);
}
}
}
}
static node_pointer expose(node_pointer cur) {
node_pointer pre = nullptr;
for (node_pointer anc = cur; anc; anc = anc->fa) {
splay(anc);
if (anc->r) {
anc->lt = dashed_tree::insert(anc->lt, ae(anc->r->fwd));
anc->r->lp = anc->lt;
}
anc->r = pre;
if (pre) {
dashed_tree::splay(pre->lp);
push(pre);
anc->lt = dashed_tree::erase(pre->lp);
}
update(anc);
pre = anc;
}
splay(cur);
return pre;
}
public:
node_pointer make_node(const Info &inf) {
node_pointer res = new node(inf);
update(res);
return res;
}
void evert(node_pointer cur) {
expose(cur);
toggle(cur);
push(cur);
}
void link(node_pointer v, node_pointer fa) {
expose(fa);
expose(v);
v->fa = fa;
fa->r = v;
update(fa);
}
void cut(node_pointer v) {
expose(v);
node_pointer fa = v->l;
assert(fa);
v->l = nullptr;
fa->fa = nullptr;
update(v);
}
void set_info(node_pointer cur, const Info &inf) {
expose(cur);
cur->inf = inf;
update(cur);
}
Info get_info(node_pointer cur) const { return cur->inf; }
Path query(node_pointer rt) {
evert(rt);
return rt->fwd;
}
Path query_subtree(node_pointer rt, node_pointer rot) {
evert(rt);
expose(rot);
node_pointer prv = rot->l;
rot->l = nullptr;
update(rot);
Path res = rot->fwd;
rot->l = prv;
update(rot);
return res;
}
};
/// @brief Index-based wrapper around `dynamic_rerooting_top_tree`.
template <class Path, class Point, class Info, auto vtx, auto cmp, auto rk,
auto ae, auto av>
class dynamic_rerooting_dp {
using top_tree_type =
dynamic_rerooting_top_tree<Path, Point, Info, vtx, cmp, rk, ae, av>;
top_tree_type tr;
std::vector<typename top_tree_type::node_pointer> nd;
public:
explicit dynamic_rerooting_dp(const std::vector<Info> &inf) : nd(inf.size()) {
for (int i = 0; i < int(inf.size()); i++) {
nd[i] = tr.make_node(inf[i]);
}
}
void link(int a, int b) {
tr.evert(nd[a]);
tr.link(nd[a], nd[b]);
}
void cut(int a, int b) {
tr.evert(nd[a]);
tr.cut(nd[b]);
}
void set_info(int idx, const Info &inf) { tr.set_info(nd[idx], inf); }
Info get_info(int idx) const { return tr.get_info(nd[idx]); }
Path query(int rt) { return tr.query(nd[rt]); }
Path query_subtree(int rt, int rot) {
return tr.query_subtree(nd[rt], nd[rot]);
}
};
} // namespace noya
namespace noya {
namespace dynamic_tree_path_affine_sum_internal {
template <class Mint> struct traits {
struct path {
Mint mul{};
Mint add{};
Mint sum{};
Mint cnt{};
};
struct point {
Mint sum{};
Mint cnt{};
};
struct info {
bool iv = true;
Mint a{};
Mint b{};
};
static path vertex(const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), x.a, Mint(1)};
}
return {x.a, x.b, Mint(0), Mint(0)};
}
static path compress(const path &fa, const path &v) {
return {fa.mul * v.mul, fa.mul * v.add + fa.add,
fa.sum + fa.mul * v.sum + fa.add * v.cnt, fa.cnt + v.cnt};
}
static point rake(const point &a, const point &b) {
return {a.sum + b.sum, a.cnt + b.cnt};
}
static point add_edge(const path &cl) { return {cl.sum, cl.cnt}; }
static path add_vertex(const point &ch, const info &x) {
if (x.iv) {
return {Mint(1), Mint(0), ch.sum + x.a, ch.cnt + Mint(1)};
}
return {x.a, x.b, ch.sum * x.a + ch.cnt * x.b, ch.cnt};
}
};
} // 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 n;
dp_type dp;
static std::vector<info>
make_info(const std::vector<Mint> &val,
const std::vector<std::pair<Mint, Mint>> &aff) {
std::vector<info> res(val.size() + aff.size());
for (int u = 0; u < int(val.size()); u++) {
res[u] = {true, val[u], Mint(0)};
}
for (int e = 0; e < int(aff.size()); e++) {
res[val.size() + e] = {false, aff[e].first, aff[e].second};
}
return res;
}
public:
dynamic_tree_path_affine_sum(const std::vector<Mint> &val,
const std::vector<std::pair<int, int>> &ee,
const std::vector<std::pair<Mint, Mint>> &aff)
: n(int(val.size())), dp(make_info(val, aff)) {
assert(n >= 1);
assert(int(ee.size()) == n - 1);
assert(aff.size() == ee.size());
for (int e = 0; e < n - 1; e++) {
auto [a, b] = ee[e];
assert(0 <= a && a < n && 0 <= b && b < n);
dp.link(n + e, a);
dp.link(n + e, b);
}
}
void set_vertex(int u, Mint va1) {
assert(0 <= u && u < n);
dp.set_info(u, {true, va1, Mint(0)});
}
void set_edge(int e, Mint mul, Mint add) {
assert(0 <= e && e < n - 1);
dp.set_info(n + e, {false, mul, add});
}
Mint query(int rt) {
assert(0 <= rt && rt < n);
return dp.query(rt).sum;
}
};
} // namespace noya