Static Top Tree (#104)

* add static-top-tree

* apply clang-format

* add doc

docs/static-top-tree-dp.md

@@ -0,0 +1,10 @@
+---
+title: Static Top Tree DP(Static Top Tree 上のDP)
+documentation_of: //graph/tree/static-top-tree-dp.hpp
+---
+
+* `vertex(int u) -> Path`: 頂点 $u$ のみからなる path cluster を生成する
+* `add_vertex(Point d, int v) -> Path`: point cluster $d$ の根に頂点 $v$ を代入して path cluster にする
+* `add_edge(Path d, int e) -> Point`: Path cluster $d$ に、辺 $e$ を追加して virtual な根を生やし、point cluster にする
+* `rake(Point l, Point r) -> Point`: point cluster $l, r$ をマージする
+* `compress(Path p, Path c, int e) -> Path`: path cluster $p, c$ ($p$ が根に近い側にある) に、辺 $e$ を使ってマージする

graph/tree/convert-rooted-tree.hpp

@@ -20,7 +20,7 @@ Graph<T> convert_rooted_tree(const Graph<T> &g, int r = 0) {
       if (v[to] == 0) {
         v[to] = 1;
         que.emplace(to);
-        rg.add_directed_edge(p, to, to.cost);
+        rg.g[p].emplace_back(to);
       }
     }
   }

graph/tree/static-top-tree-dp.hpp

@@ -0,0 +1,83 @@
+#include "static-top-tree.hpp"
+
+template <typename TreeDPInfo, typename G>
+struct StaticTopTreeDP {
+  using Path = typename TreeDPInfo::Path;
+  using Point = typename TreeDPInfo::Point;
+
+  using STT = StaticTopTree<G>;
+
+  const STT &g;
+
+  const TreeDPInfo &info;
+
+  explicit StaticTopTreeDP(const STT &g, const TreeDPInfo &info)
+      : g(g), info(info) {
+    dp.resize(g.size());
+    dfs(g.root);
+  }
+
+  Path update_vertex(int u) {
+    while (u != -1) {
+      modify(u);
+      u = g[u].p;
+    }
+    return get<Path>(dp[g.root]);
+  }
+
+  Path update_edge(int e) { return update_vertex(g.edge_to_vs[e]); }
+
+ private:
+  vector<variant<Point, Path> > dp;
+
+  void modify(int k) {
+    switch (g[k].op) {
+      case STT::Vertex:
+        dp[k] = info.vertex(k);
+        return;
+      case STT::Compress:
+        dp[k] = info.compress(get<Path>(dp[g[k].l]), get<Path>(dp[g[k].r]),
+                              g[k].e_id);
+        return;
+      case STT::Rake:
+        dp[k] = info.rake(get<Point>(dp[g[k].l]), get<Point>(dp[g[k].r]));
+        return;
+      case STT::AddEdge:
+        dp[k] = info.add_edge(get<Path>(dp[g[k].l]), g[k].e_id);
+        return;
+      case STT::AddVertex:
+        dp[k] = info.add_vertex(get<Point>(dp[g[k].l]), k);
+        return;
+    }
+  }
+
+  void dfs(int u) {
+    if (u == -1) return;
+    dfs(g[u].l);
+    dfs(g[u].r);
+    modify(u);
+  }
+};
+
+/*
+struct TreeDPInfo {
+  struct Point {
+  };
+  struct Path {
+  };
+
+  vector< int > A;
+
+  TreeDPInfo(int n): A(n) {}
+
+  Path vertex(int u) const { return Path{}; };
+
+  Path add_vertex(Point d, int u) const { return Path{}; }
+
+  Point add_edge(Path d, int e) const { return Point{}; }
+
+  Point rake(Point l, Point r) const { return Point{}; }
+
+  Path compress(Path p, Path c, int e) const { return Path{}; }
+};
+*/

graph/tree/static-top-tree.hpp

@@ -0,0 +1,139 @@
+template <typename G>
+struct StaticTopTree {
+  enum OpType { Vertex, AddVertex, AddEdge, Rake, Compress };
+
+  struct Node {
+    OpType op;
+
+    int l, r, p;
+
+    int e_id;
+
+    Node(OpType op, int l, int r) : op{op}, l{l}, r{r}, p{-1}, e_id{-1} {}
+  };
+
+  vector<Node> vs;
+
+  vector<int> edge_to_vs;
+
+  int root;
+
+  explicit StaticTopTree(G &g, int r = 0) : g(g), edge_to_vs(g.size() - 1) {
+    int e_sz = 0;
+    for (int i = 0; i < g.size(); i++) e_sz += g[i].size();
+    if (e_sz + 1 != g.size()) {
+      throw std::runtime_error("`g` must be a directed tree.");
+    }
+    dfs(r);
+
+    vs.assign(g.size(), {Vertex, -1, -1});
+    vs.reserve(g.size() * 4);
+    root = compress(r).first;
+    vs.shrink_to_fit();
+  }
+
+  const Node &operator[](int k) const { return vs[k]; }
+
+  size_t size() const { return vs.size(); }
+
+ private:
+  G &g;
+
+  using P = pair<int, int>;
+
+  int dfs(int u) {
+    int size = 1, heavy = 0;
+    for (auto &v : g[u]) {
+      int subtree_size = dfs(v);
+      size += subtree_size;
+      if (heavy < subtree_size) {
+        heavy = subtree_size;
+        swap(v, g[u][0]);
+      }
+    }
+    return size;
+  }
+
+  int make_node(OpType t, int l, int r, int k = -1) {
+    if (k == -1) {
+      k = (int)vs.size();
+      vs.emplace_back(t, l, r);
+    } else {
+      vs[k] = {t, l, r};
+    }
+    if (l != -1) vs[l].p = k;
+    if (r != -1) vs[r].p = k;
+    return k;
+  }
+
+  P merge_forRake(const vector<P> &a) {
+    if (a.size() == 1) return a[0];
+    int size_sum = 0;
+    for (auto &[_, size] : a) {
+      size_sum += size;
+    }
+    vector<P> b, c;
+    for (auto &[it, size] : a) {
+      (size_sum > size ? b : c).emplace_back(it, size);
+      size_sum -= size * 2;
+    }
+    auto [l, l_size] = merge_forRake(b);
+    auto [r, r_size] = merge_forRake(c);
+    return {make_node(Rake, l, r), l_size + r_size};
+  }
+
+  P merge_forCompress(const vector<pair<P, int>> &a) {
+    if (a.size() == 1) return a[0].first;
+    int size_sum = 0;
+    for (auto &[it, _] : a) {
+      size_sum += it.second;
+    }
+    vector<pair<P, int>> b, c;
+    for (auto &[it, _] : a) {
+      (size_sum > it.second ? b : c).emplace_back(it, _);
+      size_sum -= it.second * 2;
+    }
+    auto [l, l_size] = merge_forCompress(b);
+    auto [r, r_size] = merge_forCompress(c);
+    int idx = make_node(Compress, l, r);
+    edge_to_vs[c[0].second] = idx;
+    vs[idx].e_id = c[0].second;
+    return {idx, l_size + r_size};
+  }
+
+  P add_edge(int u, int e_idx) {
+    auto [it, size] = compress(u);
+    int idx = make_node(AddEdge, it, -1);
+    edge_to_vs[e_idx] = idx;
+    vs[idx].e_id = e_idx;
+    return {idx, size};
+  }
+
+  P rake(int u) {
+    vector<P> chs;
+    for (int j = 1; j < (int)g[u].size(); j++) {
+      chs.emplace_back(add_edge(g[u][j].to, g[u][j].idx));
+    }
+    return merge_forRake(chs);
+  }
+
+  P add_vertex(int u) {
+    if (g[u].size() < 2) {
+      return {make_node(OpType::Vertex, -1, -1, u), 1};
+    } else {
+      auto [it, size] = rake(u);
+      return {make_node(OpType::AddVertex, it, -1, u), size + 1};
+    }
+  }
+
+  P compress(int u) {
+    vector<pair<P, int>> chs{{add_vertex(u), -1}};
+    vector<int> ids{-1};
+    while (not g[u].empty()) {
+      int e_idx = g[u][0].idx;
+      u = g[u][0];
+      chs.emplace_back(add_vertex(u), e_idx);
+    }
+    return merge_forCompress(chs);
+  }
+};

test/verify/yosupo-point-set-tree-path-composite-sum-fixed-root.test.cpp

@@ -0,0 +1,73 @@
+// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum_fixed_root
+
+#include "../../template/template.hpp"
+
+#include "../../graph/graph-template.hpp"
+#include "../../graph/tree/convert-rooted-tree.hpp"
+#include "../../graph/tree/static-top-tree-dp.hpp"
+
+#include "../../math/combinatorics/montgomery-mod-int.hpp"
+
+using mint = modint998244353;
+
+struct TreeDPInfo {
+  struct Point {
+    mint val, num;
+  };
+  struct Path {
+    mint val, num, a, b;
+  };
+
+  vector< int > A, B, C;
+
+  TreeDPInfo(int n): A(n), B(n - 1), C(n - 1) {}
+
+  Path vertex(int u) const { return {A[u], 1, 1, 0}; };
+
+  Path add_vertex(Point d, int u) const { return {d.val + A[u], d.num + 1, 1, 0}; };
+
+  Point add_edge(Path d, int e) const { return {d.val * B[e] + d.num * C[e], d.num}; };
+
+  Point rake(Point l, Point r) const { return {l.val + r.val, l.num + r.num}; };
+
+  Path compress(Path p, Path c, int e) const {
+    c = {c.val * B[e] + c.num * C[e], c.num, c.a * B[e], c.b * B[e] + C[e]};
+    return {p.val + c.val * p.a + c.num * p.b, p.num + c.num, p.a * c.a, p.a * c.b + p.b};
+  };
+};
+
+int main() {
+  int N, Q;
+  cin >> N >> Q;
+  TreeDPInfo info(N);
+  for (auto &a: info.A) cin >> a;
+  Graph g(N);
+  for (int i = 0; i + 1 < N; i++) {
+    int u, v, b, c;
+    cin >> u >> v >> b >> c;
+    g.add_edge(u, v);
+    info.B[i] = b;
+    info.C[i] = c;
+  }
+  g = convert_rooted_tree(g);
+  StaticTopTree tree(g);
+  StaticTopTreeDP dp(tree, info);
+  while (Q--) {
+    int t;
+    cin >> t;
+    TreeDPInfo::Path ret;
+    if (t == 0) {
+      int w, x;
+      cin >> w >> x;
+      info.A[w] = x;
+      ret = dp.update_vertex(w);
+    } else {
+      int e, y, z;
+      cin >> e >> y >> z;
+      info.B[e] = y;
+      info.C[e] = z;
+      ret = dp.update_edge(e);
+    }
+    cout << ret.val.val() << "\n";
+  }
+}