Use this to learn the idea, then write your own version.
1#define PROBLEM "https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum"23#include "../../template/template.hpp"45#include "../../modint/montgomery-modint.hpp"67#include "../../tree/dynamic-rerooting.hpp"89using mint = LazyMontgomeryModInt<998244353>;10 11struct Path {12 mint a, b, s, x;13};14struct Point {15 mint s, x;16};17struct Info {18 bool vertex;19 mint x, y;20};21Path vertex(const Info &i) {22 if (i.vertex) return {1, 0, i.x, 1};23 return {i.x, i.y, 0, 0};24}25Path compress(const Path &p, const Path &c) {26 return {p.a * c.a, p.a * c.b + p.b, p.s + p.a * c.s + p.b * c.x, p.x + c.x};27}28Point rake(const Point &a, const Point &b) { return {a.s + b.s, a.x + b.x}; }29Point add_edge(const Path &a) { return {a.s, a.x}; }30Path add_vertex(const Point &a, const Info &i) {31 if (i.vertex) return {1, 0, a.s + i.x, a.x + 1};32 return {i.x, i.y, a.s * i.x + a.x * i.y, a.x};33}34 35using DR = DynamicRerooting<Path, Point, Info, vertex, compress, rake, add_edge,36 add_vertex>;37 38using namespace Nyaan;39 40void Nyaan::solve() {41 int N, Q;42 in(N, Q);43 vector<int> A(N);44 for (auto &x : A) in(x);45 vector<int> U(N - 1), V(N - 1), B(N - 1), C(N - 1);46 for (int i = 0; i < N - 1; i++) in(U[i], V[i], B[i], C[i]);47 48 vector<Info> info(2 * N - 1);49 for (int i = 0; i < N; i++) info[i] = {true, A[i], 0};50 for (int i = 0; i < N - 1; i++) info[N + i] = {false, B[i], C[i]};51 52 DR dr{2 * N - 1, info};53 for (int i = 0; i < N - 1; i++) {54 dr.add_edge(N + i, U[i]), dr.add_edge(N + i, V[i]);55 }56 57 while (Q--) {58 int cmd, i, x;59 in(cmd, i, x);60 if (cmd == 0) {61 dr.set_info(i, {true, x, 0});62 } else {63 int y;64 in(y);65 dr.set_info(N + i, {false, x, y});66 }67 ini(r);68 out(dr.query(r).s.get());69 }70}71