Luzhiled's Library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub ei1333/library

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

Depends on

Code

// 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";
  }
}

#line 1 "test/verify/yosupo-point-set-tree-path-composite-sum-fixed-root.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum_fixed_root

#line 1 "template/template.hpp"
#include <bits/stdc++.h>

using namespace std;

using int64 = long long;

const int64 infll = (1LL << 62) - 1;
const int inf = (1 << 30) - 1;

struct IoSetup {
  IoSetup() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
    cerr << fixed << setprecision(10);
  }
} iosetup;

template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << p.first << " " << p.second;
  return os;
}

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T>
ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++) {
    os << v[i] << (i + 1 != v.size() ? " " : "");
  }
  return os;
}

template <typename T>
istream &operator>>(istream &is, vector<T> &v) {
  for (T &in : v) is >> in;
  return is;
}

template <typename T1, typename T2>
inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}

template <typename T1, typename T2>
inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

template <typename T = int64>
vector<T> make_v(size_t a) {
  return vector<T>(a);
}

template <typename T, typename... Ts>
auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}

template <typename T, typename V>
typename enable_if<is_class<T>::value == 0>::type fill_v(T &t, const V &v) {
  t = v;
}

template <typename T, typename V>
typename enable_if<is_class<T>::value != 0>::type fill_v(T &t, const V &v) {
  for (auto &e : t) fill_v(e, v);
}

template <typename F>
struct FixPoint : F {
  explicit FixPoint(F &&f) : F(forward<F>(f)) {}

  template <typename... Args>
  decltype(auto) operator()(Args &&...args) const {
    return F::operator()(*this, forward<Args>(args)...);
  }
};

template <typename F>
inline decltype(auto) MFP(F &&f) {
  return FixPoint<F>{forward<F>(f)};
}
#line 4 "test/verify/yosupo-point-set-tree-path-composite-sum-fixed-root.test.cpp"

#line 2 "graph/graph-template.hpp"

/**
 * @brief Graph Template(グラフテンプレート)
 */
template <typename T = int>
struct Edge {
  int from, to;
  T cost;
  int idx;

  Edge() = default;

  Edge(int from, int to, T cost = 1, int idx = -1)
      : from(from), to(to), cost(cost), idx(idx) {}

  operator int() const { return to; }
};

template <typename T = int>
struct Graph {
  vector<vector<Edge<T> > > g;
  int es;

  Graph() = default;

  explicit Graph(int n) : g(n), es(0) {}

  size_t size() const { return g.size(); }

  void add_directed_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es++);
  }

  void add_edge(int from, int to, T cost = 1) {
    g[from].emplace_back(from, to, cost, es);
    g[to].emplace_back(to, from, cost, es++);
  }

  void read(int M, int padding = -1, bool weighted = false,
            bool directed = false) {
    for (int i = 0; i < M; i++) {
      int a, b;
      cin >> a >> b;
      a += padding;
      b += padding;
      T c = T(1);
      if (weighted) cin >> c;
      if (directed)
        add_directed_edge(a, b, c);
      else
        add_edge(a, b, c);
    }
  }

  inline vector<Edge<T> > &operator[](const int &k) { return g[k]; }

  inline const vector<Edge<T> > &operator[](const int &k) const { return g[k]; }
};

template <typename T = int>
using Edges = vector<Edge<T> >;
#line 2 "graph/tree/convert-rooted-tree.hpp"

#line 4 "graph/tree/convert-rooted-tree.hpp"

/**
 * @brief Convert-Rooted-Tree(根付き木に変換)
 */
template <typename T>
Graph<T> convert_rooted_tree(const Graph<T> &g, int r = 0) {
  int N = (int)g.size();
  Graph<T> rg(N);
  vector<int> v(N);
  v[r] = 1;
  queue<int> que;
  que.emplace(r);
  while (!que.empty()) {
    auto p = que.front();
    que.pop();
    for (auto &to : g[p]) {
      if (v[to] == 0) {
        v[to] = 1;
        que.emplace(to);
        rg.g[p].emplace_back(to);
      }
    }
  }
  return rg;
}
#line 1 "graph/tree/static-top-tree.hpp"
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);
  }
};
#line 2 "graph/tree/static-top-tree-dp.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 {};
  Path add_vertex(Point d, int u) const {}
  Point add_edge(Path d, int e) const {}
  Point rake(Point l, Point r) const {}
  Path compress(Path p, Path c, int e) const {}
};
*/
#line 8 "test/verify/yosupo-point-set-tree-path-composite-sum-fixed-root.test.cpp"

#line 2 "math/combinatorics/montgomery-mod-int.hpp"

template <uint32_t mod_, bool fast = false>
struct MontgomeryModInt {
 private:
  using mint = MontgomeryModInt;
  using i32 = int32_t;
  using i64 = int64_t;
  using u32 = uint32_t;
  using u64 = uint64_t;

  static constexpr u32 get_r() {
    u32 ret = mod_;
    for (i32 i = 0; i < 4; i++) ret *= 2 - mod_ * ret;
    return ret;
  }

  static constexpr u32 r = get_r();

  static constexpr u32 n2 = -u64(mod_) % mod_;

  static_assert(r * mod_ == 1, "invalid, r * mod != 1");
  static_assert(mod_ < (1 << 30), "invalid, mod >= 2 ^ 30");
  static_assert((mod_ & 1) == 1, "invalid, mod % 2 == 0");

  u32 x;

 public:
  MontgomeryModInt() : x{} {}

  MontgomeryModInt(const i64 &a)
      : x(reduce(u64(fast ? a : (a % mod() + mod())) * n2)) {}

  static constexpr u32 reduce(const u64 &b) {
    return u32(b >> 32) + mod() - u32((u64(u32(b) * r) * mod()) >> 32);
  }

  mint &operator+=(const mint &p) {
    if (i32(x += p.x - 2 * mod()) < 0) x += 2 * mod();
    return *this;
  }

  mint &operator-=(const mint &p) {
    if (i32(x -= p.x) < 0) x += 2 * mod();
    return *this;
  }

  mint &operator*=(const mint &p) {
    x = reduce(u64(x) * p.x);
    return *this;
  }

  mint &operator/=(const mint &p) {
    *this *= p.inv();
    return *this;
  }

  mint operator-() const { return mint() - *this; }

  mint operator+(const mint &p) const { return mint(*this) += p; }

  mint operator-(const mint &p) const { return mint(*this) -= p; }

  mint operator*(const mint &p) const { return mint(*this) *= p; }

  mint operator/(const mint &p) const { return mint(*this) /= p; }

  bool operator==(const mint &p) const {
    return (x >= mod() ? x - mod() : x) == (p.x >= mod() ? p.x - mod() : p.x);
  }

  bool operator!=(const mint &p) const {
    return (x >= mod() ? x - mod() : x) != (p.x >= mod() ? p.x - mod() : p.x);
  }

  u32 val() const {
    u32 ret = reduce(x);
    return ret >= mod() ? ret - mod() : ret;
  }

  mint pow(u64 n) const {
    mint ret(1), mul(*this);
    while (n > 0) {
      if (n & 1) ret *= mul;
      mul *= mul;
      n >>= 1;
    }
    return ret;
  }

  mint inv() const { return pow(mod() - 2); }

  friend ostream &operator<<(ostream &os, const mint &p) {
    return os << p.val();
  }

  friend istream &operator>>(istream &is, mint &a) {
    i64 t;
    is >> t;
    a = mint(t);
    return is;
  }

  static constexpr u32 mod() { return mod_; }
};

template <uint32_t mod>
using modint = MontgomeryModInt<mod>;
using modint998244353 = modint<998244353>;
using modint1000000007 = modint<1000000007>;
#line 10 "test/verify/yosupo-point-set-tree-path-composite-sum-fixed-root.test.cpp"

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";
  }
}

Test cases

Env Name Status Elapsed Memory
g++ example_00 :heavy_check_mark: AC 6 ms 4 MB
g++ example_01 :heavy_check_mark: AC 6 ms 4 MB
g++ n_hundreds_00 :heavy_check_mark: AC 6 ms 4 MB
g++ n_hundreds_01 :heavy_check_mark: AC 6 ms 4 MB
g++ tiny_00 :heavy_check_mark: AC 6 ms 4 MB
g++ tiny_01 :heavy_check_mark: AC 6 ms 4 MB
g++ typical_tree_max_00 :heavy_check_mark: AC 439 ms 53 MB
g++ typical_tree_max_01 :heavy_check_mark: AC 280 ms 48 MB
g++ typical_tree_max_02 :heavy_check_mark: AC 478 ms 48 MB
g++ typical_tree_max_03 :heavy_check_mark: AC 452 ms 48 MB
g++ typical_tree_max_04 :heavy_check_mark: AC 405 ms 49 MB
g++ typical_tree_max_05 :heavy_check_mark: AC 356 ms 52 MB
g++ typical_tree_max_06 :heavy_check_mark: AC 376 ms 52 MB
g++ typical_tree_max_07 :heavy_check_mark: AC 410 ms 52 MB
g++ typical_tree_max_08 :heavy_check_mark: AC 394 ms 49 MB
g++ typical_tree_max_09 :heavy_check_mark: AC 330 ms 48 MB
clang++ example_00 :heavy_check_mark: AC 6 ms 4 MB
clang++ example_01 :heavy_check_mark: AC 6 ms 4 MB
clang++ n_hundreds_00 :heavy_check_mark: AC 6 ms 4 MB
clang++ n_hundreds_01 :heavy_check_mark: AC 6 ms 4 MB
clang++ tiny_00 :heavy_check_mark: AC 6 ms 4 MB
clang++ tiny_01 :heavy_check_mark: AC 6 ms 4 MB
clang++ typical_tree_max_00 :heavy_check_mark: AC 464 ms 57 MB
clang++ typical_tree_max_01 :heavy_check_mark: AC 291 ms 48 MB
clang++ typical_tree_max_02 :heavy_check_mark: AC 501 ms 46 MB
clang++ typical_tree_max_03 :heavy_check_mark: AC 466 ms 48 MB
clang++ typical_tree_max_04 :heavy_check_mark: AC 456 ms 54 MB
clang++ typical_tree_max_05 :heavy_check_mark: AC 365 ms 55 MB
clang++ typical_tree_max_06 :heavy_check_mark: AC 402 ms 56 MB
clang++ typical_tree_max_07 :heavy_check_mark: AC 458 ms 57 MB
clang++ typical_tree_max_08 :heavy_check_mark: AC 429 ms 48 MB
clang++ typical_tree_max_09 :heavy_check_mark: AC 304 ms 50 MB
Back to top page