test/yosupo/tree/point_set_tree_path_composite_sum_fixed_root.test.cpp

• View this file on GitHub
• Last update: 2025-12-02 19:33:27+07:00
• Problem: https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum_fixed_root

Depends on

• Modint (math/modint.hpp)
• other/y_combinator.hpp
• Dynamic Tree DP (tree/dynamic_tree_dp.hpp)
• Static Top Tree (tree/static_top_tree.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum_fixed_root"

#include <bits/stdc++.h>
using namespace std;

#include "tree/static_top_tree.hpp"
#include "tree/dynamic_tree_dp.hpp"
#include "math/modint.hpp"
#include "other/y_combinator.hpp"

using mint = modint998;

struct edge {
    int to, eid;
    operator int() { return to; }
};

struct TreeDP {
    struct Info {
        mint a, b, c;
    };
    struct Point {
        mint sum, sz;
    };
    struct Path {
        mint sum, sz, b, c;
    };
    static Path vertex(Info x) { return {x.b * x.a + x.c, 1, x.b, x.c}; }
    static Path add_vertex(Point t, Info x) {
        return {x.b * (x.a + t.sum) + x.c * (1 + t.sz), 1 + t.sz, x.b, x.c};
    }
    static Point add_edge(Path t) { return {t.sum, t.sz}; }
    // how does child affect the top boundary vertex?
    static Path compress(Path p, Path c) {
        return {
            p.sum + p.b * c.sum + p.c * c.sz,
            p.sz + c.sz,
            p.b * c.b,
            p.b * c.c + p.c
        };
    }
    static Point rake(Point l, Point r) { return {l.sum + r.sum, l.sz + r.sz}; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<basic_string<edge>> g(n);
    vector<mint> a(n), b(n), c(n);
    for (mint &i : a) cin >> i;
    for (int i = 0, u, v; i < n - 1; i++) {
        cin >> u >> v >> b[i] >> c[i];
        g[u] += {v, i}, g[v] += {u, i};
    }
    b[n - 1] = 1, c[n - 1] = 0;

    vector<int> v_e(n, n - 1), e_v(n, n - 1);
    y_comb([&](auto &dfs, int v, int p) -> void {
        auto it = find(begin(g[v]), end(g[v]), p);
        if (it != end(g[v])) g[v].erase(it);
        for (auto [u, eid] : g[v]) v_e[u] = eid, e_v[eid] = u, dfs(u, v);
    })(0, -1);

    STT stt(g);
    DDP ddp(n, stt, TreeDP());
    for (int i = 0; i < n; i++) {
        int e = v_e[i];
        ddp.info[i] = {a[i], b[e], c[e]};
    }
    ddp.build();

    while (q--) {
        int op, x, y, z;
        cin >> op;
        if (op == 0) {
            cin >> x >> y;
            ddp.info[x].a = y;
            ddp.pull_from(x);
        } else {
            cin >> x >> y >> z;
            int v = e_v[x];
            ddp.info[v].b = y, ddp.info[v].c = z;
            ddp.pull_from(v);
        }
        cout << ddp.query().sum << "\n";
    }
}

#line 1 "test/yosupo/tree/point_set_tree_path_composite_sum_fixed_root.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum_fixed_root"

#include <bits/stdc++.h>
using namespace std;

#line 2 "tree/static_top_tree.hpp"

#line 2 "other/y_combinator.hpp"

template<class F> struct y_comb_t : F {
    template<class T> y_comb_t(T &&_f) : F(forward<T>(_f)) {}
    template<class... Args> decltype(auto) operator()(Args &&...args) {
        return F::operator()(/* ref */(*this), forward<Args>(args)...);
    }
};
template<class F> y_comb_t<decay_t<F>> y_comb(F &&f) { return {forward<F>(f)}; }
#line 4 "tree/static_top_tree.hpp"

enum Type : char { Vertex, AddVertex, AddEdge, Compress, Rake };
template<class G> struct STT {
    struct node {
        Type t;
        int p = -1, l, r;

        node() {}
        node(Type tp, int _l, int _r) : t(tp), l(_l), r(_r) {}
    };

    int rt;
    G &g;
    vector<node> vs;

    STT(G &_g, int root = 0) : g(_g) {
        int n = (int)g.size();
        vs.assign(n, {Vertex, -1, -1});
        vs.reserve(4 * n);
        y_comb([&](auto &hld, int v) -> int {
            int s = 1, mx = 0;
            for (auto &c : g[v]) {
                int sc = hld(c);
                s += sc;
                if (sc > mx) mx = sc, swap(g[v][0], c);
            }
            return s;
        })(root);
        rt = compress(root)[0];
        vs.shrink_to_fit();
    }

    int size() const { return (int)vs.size(); }
    auto &operator[](int i) const { return vs[i]; }

    using P = array<int, 2>;
    int add(int k, int l, int r, Type t) {
        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(const vector<P> &a, Type t) {
        if (a.size() == 1) return a[0];
        int s = 0;
        for (auto [v, sz] : a) s += sz;
        vector<P> b, c;
        for (auto x : a) (s > x[1] ? b : c).push_back(x), s -= 2 * x[1];
        auto [l, szl] = merge(b, t);
        auto [r, szr] = merge(c, t);
        return {add(-1, l, r, t), szl + szr};
    }
    P compress(int v) {
        vector<P> ch{add_vertex(v)};
        while (g[v].size()) ch.push_back(add_vertex(v = g[v][0]));
        return merge(ch, Compress);
    }
    P rake(int v) {
        vector<P> ch;
        for (int i = 1; i < g[v].size(); i++) ch.push_back(add_edge(g[v][i]));
        return ch.size() ? merge(ch, Rake) : P{-1, 0};
    }
    P add_edge(int v) {
        auto [u, sz] = compress(v);
        return {add(-1, u, -1, AddEdge), sz};
    }
    P add_vertex(int v) {
        auto [u, sz] = rake(v);
        return {add(v, u, -1, u == -1 ? Vertex : AddVertex), sz + 1};
    }
};
#line 2 "tree/dynamic_tree_dp.hpp"

template<class TreeDP, class STT> struct DDP {
    using Info = typename TreeDP::Info;
    using Point = typename TreeDP::Point;
    using Path = typename TreeDP::Path;

    STT &stt;
    int n, m, rt;
    vector<Info> info;
    vector<Point> point;
    vector<Path> path;

    void pull(int v) {
        if (stt[v].t == Vertex) {
            path[v] = TreeDP::vertex(info[v]);
        } else if (stt[v].t == AddVertex) {
            path[v] = TreeDP::add_vertex(point[stt[v].l], info[v]);
        } else if (stt[v].t == AddEdge) {
            point[v] = TreeDP::add_edge(path[stt[v].l]);
        } else if (stt[v].t == Compress) {
            path[v] = TreeDP::compress(path[stt[v].l], path[stt[v].r]);
        } else if (stt[v].t == Rake) {
            point[v] = TreeDP::rake(point[stt[v].l], point[stt[v].r]);
        }
    }

    DDP(int _n, STT &_stt, [[maybe_unused]] TreeDP _dp) :
        stt(_stt), n(_n), m(stt.size()), info(n), point(m), path(m) {}
    void build() {
        y_comb([&](auto &dfs, int v) -> void {
            if (stt[v].l != -1) dfs(stt[v].l);
            if (stt[v].r != -1) dfs(stt[v].r);
            pull(v);
        })(stt.rt);
    }
    void pull_from(int v) {
        for (; v != -1; v = stt[v].p) pull(v);
    }
    void set(int v, const Info &x) { info[v] = x, pull_from(v); }
    Path query() { return path[stt.rt]; }
};
#line 2 "math/modint.hpp"

// clang-format off
template<uint32_t m> struct modint {
    static_assert(m >= 1, "Modulus must be in the range [1;2^31)");

    using mint = modint;
    static constexpr bool is_simple = true;

    static constexpr uint32_t mod() noexcept { return m; }
    constexpr modint() noexcept = default;
    constexpr modint(int64_t v) noexcept : _v(uint32_t((v %= m) < 0 ? v + m : v)) {}
    constexpr static mint raw(uint32_t v) noexcept { mint x; return x._v = v, x; }
    template<class T> constexpr explicit operator T() const noexcept { return _v; }

    constexpr mint &operator++() noexcept { return _v = ++_v == mod() ? 0 : _v, *this; }
    constexpr mint &operator--() noexcept { --(_v ? _v : _v = mod()); return *this; }
    constexpr mint operator++(int) noexcept { return exchange(*this, ++mint(*this)); }
    constexpr mint operator--(int) noexcept { return exchange(*this, --mint(*this)); }

    constexpr mint &operator+=(mint rhs) noexcept {
        return _v = int(_v += rhs._v - mod()) < 0 ? _v + mod() : _v, *this;
    }
    constexpr mint &operator-=(mint rhs) noexcept {
        return _v = int(_v -= rhs._v) < 0 ? _v + mod() : _v, *this;
    }
    constexpr mint &operator*=(mint rhs) noexcept {
        return _v = uint64_t(_v) * rhs._v % mod(), *this;
    }
    constexpr mint &operator/=(mint rhs) noexcept {
        return *this = *this * rhs.inv();
    }

    constexpr friend mint operator+(mint l, mint r) noexcept { return l += r; }
    constexpr friend mint operator-(mint l, mint r) noexcept { return l -= r; }
    constexpr friend mint operator*(mint l, mint r) noexcept { return l *= r; }
    constexpr friend mint operator/(mint l, mint r) noexcept { return l /= r; }

    constexpr mint operator+() const noexcept { return *this; }
    constexpr mint operator-() const noexcept { return raw(_v ? mod() - _v : 0); }

    constexpr friend bool operator==(mint l, mint r) noexcept { return l._v == r._v; }
    constexpr friend bool operator!=(mint l, mint r) noexcept { return l._v != r._v; }
    constexpr friend bool operator<(mint l, mint r) noexcept { return l._v < r._v; }

    constexpr mint pow(uint64_t n) const noexcept {
        mint b = *this, res = 1;
        while (n) n & 1 ? res *= b : 0, b *= b, n >>= 1;
        return res;
    }

    constexpr mint inv() const noexcept {
        int a = _v, b = mod(), x = 1, y = 0;
        while (b) {
            x = exchange(y, x - a / b * y);
            a = exchange(b, a % b);
        }
        assert(a == 1);
        return x;
    }

    friend istream &operator>>(istream &is, mint &x) {
        int64_t v{};
        return is >> v, x = v, is;
    }
    friend ostream &operator<<(ostream &os, const mint &x) { return os << x._v; }

  private:
    uint32_t _v = 0;
};
using modint107 = modint<1'000'000'007>;
using modint998 = modint<998'244'353>;
// clang-format on
#line 10 "test/yosupo/tree/point_set_tree_path_composite_sum_fixed_root.test.cpp"

using mint = modint998;

struct edge {
    int to, eid;
    operator int() { return to; }
};

struct TreeDP {
    struct Info {
        mint a, b, c;
    };
    struct Point {
        mint sum, sz;
    };
    struct Path {
        mint sum, sz, b, c;
    };
    static Path vertex(Info x) { return {x.b * x.a + x.c, 1, x.b, x.c}; }
    static Path add_vertex(Point t, Info x) {
        return {x.b * (x.a + t.sum) + x.c * (1 + t.sz), 1 + t.sz, x.b, x.c};
    }
    static Point add_edge(Path t) { return {t.sum, t.sz}; }
    // how does child affect the top boundary vertex?
    static Path compress(Path p, Path c) {
        return {
            p.sum + p.b * c.sum + p.c * c.sz,
            p.sz + c.sz,
            p.b * c.b,
            p.b * c.c + p.c
        };
    }
    static Point rake(Point l, Point r) { return {l.sum + r.sum, l.sz + r.sz}; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<basic_string<edge>> g(n);
    vector<mint> a(n), b(n), c(n);
    for (mint &i : a) cin >> i;
    for (int i = 0, u, v; i < n - 1; i++) {
        cin >> u >> v >> b[i] >> c[i];
        g[u] += {v, i}, g[v] += {u, i};
    }
    b[n - 1] = 1, c[n - 1] = 0;

    vector<int> v_e(n, n - 1), e_v(n, n - 1);
    y_comb([&](auto &dfs, int v, int p) -> void {
        auto it = find(begin(g[v]), end(g[v]), p);
        if (it != end(g[v])) g[v].erase(it);
        for (auto [u, eid] : g[v]) v_e[u] = eid, e_v[eid] = u, dfs(u, v);
    })(0, -1);

    STT stt(g);
    DDP ddp(n, stt, TreeDP());
    for (int i = 0; i < n; i++) {
        int e = v_e[i];
        ddp.info[i] = {a[i], b[e], c[e]};
    }
    ddp.build();

    while (q--) {
        int op, x, y, z;
        cin >> op;
        if (op == 0) {
            cin >> x >> y;
            ddp.info[x].a = y;
            ddp.pull_from(x);
        } else {
            cin >> x >> y >> z;
            int v = e_v[x];
            ddp.info[v].b = y, ddp.info[v].c = z;
            ddp.pull_from(v);
        }
        cout << ddp.query().sum << "\n";
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	4 MB
g++	example_01	AC	4 ms	4 MB
g++	n_hundreds_00	AC	5 ms	4 MB
g++	n_hundreds_01	AC	5 ms	4 MB
g++	tiny_00	AC	4 ms	3 MB
g++	tiny_01	AC	4 ms	4 MB
g++	typical_tree_max_00	AC	356 ms	51 MB
g++	typical_tree_max_01	AC	273 ms	44 MB
g++	typical_tree_max_02	AC	359 ms	37 MB
g++	typical_tree_max_03	AC	366 ms	38 MB
g++	typical_tree_max_04	AC	327 ms	44 MB
g++	typical_tree_max_05	AC	296 ms	47 MB
g++	typical_tree_max_06	AC	315 ms	47 MB
g++	typical_tree_max_07	AC	348 ms	48 MB
g++	typical_tree_max_08	AC	348 ms	37 MB
g++	typical_tree_max_09	AC	282 ms	40 MB
g++	typical_tree_max_degree_00	AC	364 ms	51 MB
g++	typical_tree_max_degree_01	AC	240 ms	44 MB
g++	typical_tree_max_degree_02	AC	366 ms	37 MB
g++	typical_tree_max_degree_03	AC	328 ms	38 MB
g++	typical_tree_max_degree_04	AC	348 ms	44 MB
g++	typical_tree_max_degree_05	AC	286 ms	47 MB
g++	typical_tree_max_degree_06	AC	322 ms	47 MB
g++	typical_tree_max_degree_07	AC	377 ms	48 MB
g++	typical_tree_max_degree_08	AC	328 ms	37 MB
g++	typical_tree_max_degree_09	AC	264 ms	40 MB
clang++	example_00	AC	5 ms	4 MB
clang++	example_01	AC	4 ms	4 MB
clang++	n_hundreds_00	AC	4 ms	4 MB
clang++	n_hundreds_01	AC	4 ms	4 MB
clang++	tiny_00	AC	4 ms	4 MB
clang++	tiny_01	AC	4 ms	4 MB
clang++	typical_tree_max_00	AC	355 ms	51 MB
clang++	typical_tree_max_01	AC	262 ms	44 MB
clang++	typical_tree_max_02	AC	396 ms	37 MB
clang++	typical_tree_max_03	AC	331 ms	38 MB
clang++	typical_tree_max_04	AC	373 ms	44 MB
clang++	typical_tree_max_05	AC	318 ms	47 MB
clang++	typical_tree_max_06	AC	335 ms	47 MB
clang++	typical_tree_max_07	AC	340 ms	48 MB
clang++	typical_tree_max_08	AC	354 ms	37 MB
clang++	typical_tree_max_09	AC	298 ms	40 MB
clang++	typical_tree_max_degree_00	AC	359 ms	51 MB
clang++	typical_tree_max_degree_01	AC	246 ms	44 MB
clang++	typical_tree_max_degree_02	AC	383 ms	37 MB
clang++	typical_tree_max_degree_03	AC	366 ms	38 MB
clang++	typical_tree_max_degree_04	AC	329 ms	44 MB
clang++	typical_tree_max_degree_05	AC	316 ms	47 MB
clang++	typical_tree_max_degree_06	AC	339 ms	47 MB
clang++	typical_tree_max_degree_07	AC	362 ms	48 MB
clang++	typical_tree_max_degree_08	AC	357 ms	38 MB
clang++	typical_tree_max_degree_09	AC	289 ms	40 MB

Back to top page