test/yosupo/tree/point_set_tree_path_composite_sum_top_tree.test.cpp

• View this file on GitHub
• Last update: 2025-09-26 12:32:09+07:00
• Problem: https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum

Depends on

• Modint (math/modint.hpp)
• Top Tree (tree/top_tree.hpp)

Code

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

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

#include "math/modint.hpp"
#include "tree/top_tree.hpp"

using mint = modint998;

struct TreeDP {
    struct Info {
        bool is_vertex;
        mint a, b;
    };
    struct Point {
        mint sum, sz;
    };
    struct Path {
        mint sum, sz, b, c;
    };
    static Path vertex(Info x) {
        if (x.is_vertex) return {x.a, 1, 1, 0};
        return {0, 0, x.a, x.b};
    }
    static Path add_vertex(Point t, Info x) {
        if (x.is_vertex) return {t.sum + x.a, t.sz + 1, 1, 0};
        return {x.a * t.sum + x.b * t.sz, t.sz, x.a, x.b};
    }
    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;
    TopTree<TreeDP> tr(2 * n - 1);
    for (int i = 0, a; i < n; i++) cin >> a, tr.set(i, {true, a, 0});
    for (int i = 0, u, v, b, c; i < n - 1; i++) {
        cin >> u >> v >> b >> c;
        tr.set(i + n, {false, b, c});
        tr.link(u, i + n);
        tr.link(v, i + n);
    }
    while (q--) {
        int op, r;
        cin >> op;
        if (op == 0) {
            int w, x;
            cin >> w >> x >> r;
            tr.set(w, {true, x, 0});
        } else {
            int e, y, z;
            cin >> e >> y >> z >> r;
            tr.set(e + n, {false, y, z});
        }
        cout << tr.fold(r).sum << "\n";
    }
}

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

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

#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 2 "tree/top_tree.hpp"

// reference:
// https://ei1333.github.io/library/structure/dynamic-tree/top-tree.hpp
template<class TreeDP> struct TopTree {
    using Info = typename TreeDP::Info;
    using Point = typename TreeDP::Point;
    using Path = typename TreeDP::Path;

    struct node;
    using ptr = node *;
    vector<node> nodes;

    TopTree(int n = 0) : nodes(n) {}
    TopTree(const vector<Info> &v) : nodes(v.size()) {
        for (int i = 0; i < v.size(); i++) {
            nodes[i].info = v[i];
            nodes[i].pull();
        }
    }

    auto operator[](int i) { return &nodes[i]; }
    int id(ptr t) { return int(t - &nodes[0]); }

    void splay(int v) { splay(&nodes[v]); }
    void expose(int v) { expose(&nodes[v]); }
    void evert(int v) { evert(&nodes[v]); }
    void link(int v, int p) { link(&nodes[v], &nodes[p]); }
    void cut(int v) { cut(&nodes[v]); }
    void cut(int v, int r) { evert(r), cut(v); }
    void set(int v, const Info &x) {
        apply(v, [&x](auto &info) { info = x; });
    }
    template<class F> void apply(int v, F &&f) {
        expose(v), f(nodes[v].info), nodes[v].pull();
    }
    int lca(int u, int v) {
        ptr l = lca(&nodes[u], &nodes[v]);
        return l ? id(l) : -1;
    }
    Path fold(int v) { return evert(v), nodes[v].sum; }
    Path fold_path(int v) { return expose(v), nodes[v].sum; }
    Path fold_path(int u, int v) { return evert(u), fold_path(v); }
    Path fold_subtree(int v) {
        expose(v);
        ptr l = nodes[v].l;
        nodes[v].l = 0, nodes[v].pull();
        Path ret = nodes[v].sum;
        nodes[v].l = l, nodes[v].pull();
        return ret;
    }
    Path fold_subtree(int v, int r) { return evert(r), fold_subtree(v); }

    struct rake_node {
        using rptr = rake_node *;
        rptr p = 0, l = 0, r = 0;
        Point val{}, sum{};

        rake_node(const Point &x) : val(x), sum(x) {}

        void pull() {
            sum = val;
            if (l) sum = TreeDP::rake(sum, l->sum);
            if (r) sum = TreeDP::rake(sum, r->sum);
        }

        rptr &ch(bool d) { return !d ? l : r; }
        static bool dir(rptr t) { return t == t->p->r; }
        static void rot(rptr t) {
            rptr p = t->p;
            int d = dir(t);
            t->p = p->p;
            if (p->p) p->p->ch(dir(p)) = t;
            if ((p->ch(d) = t->ch(!d))) p->ch(d)->p = p;
            t->ch(!d) = p, p->p = t;
            p->pull(), t->pull();
        }
        static void splay(rptr t) {
            for (; t->p; rot(t))
                if (t->p->p) dir(t) == dir(t->p) ? rot(t->p) : rot(t);
        }

        rptr rightmost() {
            rptr cur = this;
            while (cur->r) cur = cur->r;
            return cur;
        }
        friend rptr insert(rptr t, const Point &x) {
            rptr z = new rake_node(x);
            if (!t) return z;
            splay(t = t->rightmost());
            t->r = z, z->p = t;
            t->pull();
            return splay(z), z;
        }
        friend rptr erase(rptr t) {
            splay(t);
            rptr l = t->l, r = t->r;
            delete t;
            if (l) l->p = 0;
            if (r) r->p = 0;
            if (!l) {
                t = r;
            } else if (!r) {
                t = l;
            } else {
                splay(t = l->rightmost());
                t->r = r, r->p = t;
                t->pull();
            }
            return t;
        }
    };
    struct node {
        ptr p = 0, l = 0, r = 0;
        rake_node *light = 0, *point = 0;
        bool rev = 0;
        Info info{};
        Path sum{}, mus{};

        ~node() {
            auto del = [](auto f, auto t) {
                if (!t) return;
                f(f, t->l), f(f, t->r);
                delete t;
            };
            del(del, light);
        }


        void pull() {
            Path raked = light ? TreeDP::add_vertex(light->sum, info)
                                : TreeDP::vertex(info);
            sum = mus = raked;
            if (l) {
                sum = TreeDP::compress(l->sum, sum);
                mus = TreeDP::compress(mus, l->mus);
            }
            if (r) {
                sum = TreeDP::compress(sum, r->sum);
                mus = TreeDP::compress(r->mus, mus);
            }
        }
        void flip() { swap(l, r), swap(sum, mus), rev ^= 1; }
        void push() {
            if (exchange(rev, 0)) {
                if (l) l->flip();
                if (r) r->flip();
            }
        }

        ptr &ch(bool d) { return !d ? l : r; }
        static bool is_root(ptr t) {
            return !t->p || (t != t->p->l && t != t->p->r);
        }
        static bool dir(ptr t) { return t == t->p->r; }
        static void rot(ptr t) {
            ptr p = t->p;
            int d = dir(t);
            t->p = p->p;
            if (!is_root(p)) p->p->ch(dir(p)) = t;
            if ((p->ch(d) = t->ch(!d))) p->ch(d)->p = p;
            t->ch(!d) = p, p->p = t;
            p->pull(), t->pull();
        }
        static void splay(ptr t) {
            t->push();
            {
                ptr cur = t;
                while (!is_root(cur)) cur = cur->p;
                t->point = cur->point;
                if (t != cur) cur->point = 0;
            }
            for (; !is_root(t); rot(t)) {
                ptr p = t->p, g = p->p;
                if (g) g->push();
                p->push(), t->push();
                if (!is_root(p)) rot(dir(t) == dir(p) ? p : t);
            }
        }
    };

    static ptr expose(ptr t) {
        ptr prv = 0;
        for (ptr cur = t; cur; cur = cur->p) {
            node::splay(cur);
            if (cur->r) {
                cur->light = insert(cur->light, TreeDP::add_edge(cur->r->sum));
                cur->r->point = cur->light;
            }
            if (prv) {
                prv->push();
                cur->light = erase(prv->point);
            }
            cur->r = exchange(prv, cur), cur->pull();
        }
        node::splay(t);
        return prv;
    }
    static void evert(ptr t) { expose(t), t->flip(); }
    static void link(ptr v, ptr p) {
        evert(v), expose(p);
        assert(!p->r);
        p->r = v, v->p = p, p->pull();
    }
    static void cut(ptr v) {
        expose(v);
        assert(v->l);
        v->l = v->l->p = 0, v->pull();
    }
    static ptr lca(ptr u, ptr v) {
        if (u == v) return u;
        expose(u);
        ptr l = expose(v);
        return u->p ? l : 0;
    }
};
#line 8 "test/yosupo/tree/point_set_tree_path_composite_sum_top_tree.test.cpp"

using mint = modint998;

struct TreeDP {
    struct Info {
        bool is_vertex;
        mint a, b;
    };
    struct Point {
        mint sum, sz;
    };
    struct Path {
        mint sum, sz, b, c;
    };
    static Path vertex(Info x) {
        if (x.is_vertex) return {x.a, 1, 1, 0};
        return {0, 0, x.a, x.b};
    }
    static Path add_vertex(Point t, Info x) {
        if (x.is_vertex) return {t.sum + x.a, t.sz + 1, 1, 0};
        return {x.a * t.sum + x.b * t.sz, t.sz, x.a, x.b};
    }
    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;
    TopTree<TreeDP> tr(2 * n - 1);
    for (int i = 0, a; i < n; i++) cin >> a, tr.set(i, {true, a, 0});
    for (int i = 0, u, v, b, c; i < n - 1; i++) {
        cin >> u >> v >> b >> c;
        tr.set(i + n, {false, b, c});
        tr.link(u, i + n);
        tr.link(v, i + n);
    }
    while (q--) {
        int op, r;
        cin >> op;
        if (op == 0) {
            int w, x;
            cin >> w >> x >> r;
            tr.set(w, {true, x, 0});
        } else {
            int e, y, z;
            cin >> e >> y >> z >> r;
            tr.set(e + n, {false, y, z});
        }
        cout << tr.fold(r).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	4 MB
g++	tiny_01	AC	4 ms	4 MB
g++	typical_tree_max_00	AC	1204 ms	42 MB
g++	typical_tree_max_01	AC	722 ms	57 MB
g++	typical_tree_max_02	AC	1460 ms	48 MB
g++	typical_tree_max_03	AC	1276 ms	43 MB
g++	typical_tree_max_04	AC	1279 ms	47 MB
g++	typical_tree_max_05	AC	933 ms	49 MB
g++	typical_tree_max_06	AC	1041 ms	47 MB
g++	typical_tree_max_07	AC	1190 ms	43 MB
g++	typical_tree_max_08	AC	1327 ms	51 MB
g++	typical_tree_max_09	AC	806 ms	49 MB
g++	typical_tree_max_degree_00	AC	1175 ms	43 MB
g++	typical_tree_max_degree_01	AC	518 ms	57 MB
g++	typical_tree_max_degree_02	AC	1397 ms	48 MB
g++	typical_tree_max_degree_03	AC	1256 ms	43 MB
g++	typical_tree_max_degree_04	AC	1244 ms	47 MB
g++	typical_tree_max_degree_05	AC	822 ms	49 MB
g++	typical_tree_max_degree_06	AC	963 ms	47 MB
g++	typical_tree_max_degree_07	AC	1196 ms	43 MB
g++	typical_tree_max_degree_08	AC	1271 ms	51 MB
g++	typical_tree_max_degree_09	AC	712 ms	49 MB
clang++	example_00	AC	4 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	1201 ms	43 MB
clang++	typical_tree_max_01	AC	791 ms	56 MB
clang++	typical_tree_max_02	AC	1572 ms	48 MB
clang++	typical_tree_max_03	AC	1366 ms	43 MB
clang++	typical_tree_max_04	AC	1332 ms	47 MB
clang++	typical_tree_max_05	AC	997 ms	49 MB
clang++	typical_tree_max_06	AC	1077 ms	47 MB
clang++	typical_tree_max_07	AC	1238 ms	43 MB
clang++	typical_tree_max_08	AC	1443 ms	51 MB
clang++	typical_tree_max_09	AC	879 ms	49 MB
clang++	typical_tree_max_degree_00	AC	1265 ms	43 MB
clang++	typical_tree_max_degree_01	AC	562 ms	56 MB
clang++	typical_tree_max_degree_02	AC	1525 ms	48 MB
clang++	typical_tree_max_degree_03	AC	1291 ms	43 MB
clang++	typical_tree_max_degree_04	AC	1286 ms	47 MB
clang++	typical_tree_max_degree_05	AC	856 ms	49 MB
clang++	typical_tree_max_degree_06	AC	985 ms	47 MB
clang++	typical_tree_max_degree_07	AC	1233 ms	43 MB
clang++	typical_tree_max_degree_08	AC	1395 ms	51 MB
clang++	typical_tree_max_degree_09	AC	787 ms	49 MB

Back to top page