test/yosupo/tree/point_set_tree_path_composite_sum_top_tree_ecnerwala.test.cpp

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

Depends on

• Modint (math/modint.hpp)
• Top Tree (ecnerwala) (tree/top_tree_ecnerwala.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_ecnerwala.hpp"

using mint = modint998;
struct affine {
    mint a = 1, b = 0;
    affine operator()(const affine &g) const { return {a * g.a, a * g.b + b}; }
    mint operator()(mint x) const { return a * x + b; }
};

struct node : top_tree_node_base<node> {
    mint val;
    affine f;
    int subtree_size;
    array<affine, 2> path_affine;
    array<mint, 2> ans;

    void flip() { swap(path_affine[0], path_affine[1]), swap(ans[0], ans[1]); }
    void pull() {
        subtree_size = is_vert;
        path_affine.fill({});
        ans.fill(0);
        for (int i = 0; i <= 2; i++) {
            if (!c[i]) continue;
            subtree_size += c[i]->subtree_size;
        }
        if (is_vert) { // add_vertex(self, rake(c[0], c[1]))
            ans[0] = val;
            for (int i = 0; i < 2; i++) {
                if (!c[i]) continue;
                ans[0] += c[i]->ans[0];
            }
            ans[1] = ans[0];
        } else if (is_path) { // compress(c[0], self, c[1])
            for (int i = 0; i < 2; i++) {
                affine tmp = c[i]->path_affine[i](f);
                path_affine[i] = tmp(c[!i]->path_affine[i]);
                ans[i] = c[i]->ans[i] + tmp.a * c[!i]->ans[i] +
                         tmp.b * c[!i]->subtree_size;
            }
        } else { // rake(c[0], add_edge(self, c[2]), c[1])
            for (int i = 0; i < 2; i++) {
                if (!c[i]) continue;
                ans[0] += c[i]->ans[0];
            }
            ans[0] += f.a * c[2]->ans[0] + f.b * c[2]->subtree_size;
            ans[1] = ans[0];
        }
    }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<node> t(2 * n - 1);
    for (int i = 0; i < n; i++) {
        cin >> t[i].val;
        t[i].set_vert();
    }
    for (int i = 0, u, v; i < n - 1; i++) {
        cin >> u >> v >> t[i + n].f.a >> t[i + n].f.b;
        link(t[i + n], t[u], t[v]);
    }
    for (int qi = 0; qi < q; qi++) {
        int op, r;
        cin >> op;
        if (op == 0) {
            int w, x;
            cin >> w >> x;
            t[w].expose();
            t[w].val = x;
            t[w].pull_all();
        } else if (op == 1) {
            int e, y, z;
            cin >> e >> y >> z;
            auto &edge = t[e + n];
            edge.expose_edge();
            edge.f = {y, z};
            edge.pull_all();
        }
        cin >> r;
        t[r].evert();
        cout << t[r].ans[0] << "\n";
    }
}

#line 1 "test/yosupo/tree/point_set_tree_path_composite_sum_top_tree_ecnerwala.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_ecnerwala.hpp"

/**
 * https://github.com/ecnerwala/cp-book/blob/master/src/top_tree.hpp
 * Convenient memo:
 *   Creating vertices:
 *     n->val = ...;
 *     n->set_vert();
 *
 *   Creating edges:
 *     link(e, va, vb);
 *
 *   Updates:
 *     auto cur = get_path(va, vb); // or get_subtree(va, vb)
 *     cur->do_stuff();
 *     cur->pull_all();
 *
 * Node types:
 *   path edges: compress(c[0], self, c[1])
 *     assert(is_path && !is_vert);
 *     assert(c[0] && c[1]);
 *     assert(c[0]->is_path && c[1]->is_path);
 *     assert(!c[2]);
 *   (path) vertices: add_vertex(self, rake(c[0], c[1]))
 *     assert(is_path && is_vert);
 *     assert(!c[2]);
 *     if (c[0]) assert(!c[0]->is_path);
 *     if (c[1]) assert(!c[1]->is_path);
 *   non-path edges: rake(c[0], add_edge(self, c[2]), c[1])
 *     assert(!is_path && !is_vert);
 *     assert(c[2])
 *     assert(c[2]->is_path);
 *     if (c[0]) assert(!c[0]->is_path);
 *     if (c[1]) assert(!c[1]->is_path);
 */
template<class node> struct top_tree_node_base {
    using ptr = node *;

  private:
    ptr as_derived() { return (ptr)this; }
    auto as_derived() const { return (const ptr)this; }

  public:
    bool rev = false;
    ptr p = 0;
    array<ptr, 3> c{0, 0, 0};

    int d() const {
        assert(p);
        if (this == p->c[0]) return 0;
        if (this == p->c[1]) return 1;
        if (this == p->c[2]) return 2;
        assert(false);
    }
    ptr &p_c() { return p->c[d()]; }

    bool is_vert, is_path;
    void set_vert() { is_path = is_vert = true, _pull(); }

    bool r() const { return !p || p->is_path != is_path; }

  protected:
#define CHECK(a) \
    template<class T, class = void> struct has_##a##_t : false_type {}; \
    template<class T> struct has_##a##_t<T, void_t<decltype(&T::a)>> : true_type {}; \
    template<class T> static constexpr bool has_##a = has_##a##_t<T>::value;
    CHECK(push) CHECK(flip) CHECK(pull)
#undef CHECK
    void _flip() {
        assert(is_path);
        swap(c[0], c[1]), rev ^= 1;
        if constexpr (has_flip<node>) as_derived()->flip();
    }
    void push_rev() {
        if (rev) {
            assert(is_path);
            if (!is_vert) {
                c[0]->_flip();
                c[1]->_flip();
            }
            rev = false;
        }
    }
    void _push() {
        push_rev();
        if constexpr (has_push<node>) as_derived()->push();
    }
    void _pull() {
        if constexpr (has_pull<node>) as_derived()->pull();
    }

  public:
    void push_all() {
        if (p) p->push_all();
        _push();
    }
    void push_flip_all() {
        if (p) p->push_flip_all();
        if (rev) {
            assert(is_path);
            if (!is_vert) {
                c[0]->_flip();
                c[1]->_flip();
            }
            rev = false;
        }
    }
    ptr pull_all() {
        ptr cur = as_derived();
        cur->_push(), cur->_pull();
        for (; cur->p; cur->_pull()) cur = cur->p;
        return cur;
    }

  private:
    static inline void attach(ptr pa, int c_d, ptr ch) {
        pa->c[c_d] = ch;
        if (ch) ch->p = pa;
    }
    void rot() {
        assert(!is_vert);
        assert(!r());

        ptr pa = p;
        int x = d();
        assert(x != 2);
        ptr ch = c[!x];

        if (pa->p) pa->p_c() = as_derived();
        this->p = pa->p;

        attach(pa, x, ch);
        attach(as_derived(), !x, pa);

        pa->_pull();
    }
    void rot2(int c_d) {
        assert(!is_vert);
        assert(!r());
        assert(c[c_d]);
        assert(!c[c_d]->is_vert);

        if (d() == c_d) return rot();

        ptr pa = p;
        int x = d();
        assert(x != 2);
        ptr ch = c[c_d]->c[!x];

        if (pa->p) pa->p_c() = as_derived();
        this->p = pa->p;

        attach(pa, x, ch);
        attach(c[c_d], !x, pa);

        pa->_pull();
    }

    void splay_dir(int x) {
        for (; !r() && d() == x; rot())
            if (!p->r() && p->d() == x) p->rot();
    }

    void splay() {
        assert(!is_vert);
        for (; !r(); rot())
            if (!p->r()) p->d() == d() ? p->rot() : rot();
    }

    void splay2(int c_d) {
        assert(!is_vert && is_path);
        assert(c[c_d] && !c[c_d]->is_vert);
        for (; !r(); rot2(c_d))
            if (!p->r()) p->d() == d() ? p->rot() : rot2(c_d);
    }

    void splay2() {
        assert(!is_vert && is_path);
        assert(!r());
        p->splay2(d());
    }

    void splay_vert() {
        assert(is_vert);

        if (r()) return;

        p->splay_dir(d());
        if (p->r()) return;
        assert(p->d() != d());

        if (d() == 1) p->rot();
        assert(d() == 0);

        p->splay2();
        assert(d() == 0);
        assert(p->d() == 1);
        assert(p->p->r());
    }

    ptr cut_right() {
        assert(is_vert && is_path);
        splay_vert();

        if (r() || d() == 1) {
            assert(r() || (d() == 1 && p->r()));
            assert(!c[0]);
            return 0;
        }

        ptr pa = p;
        assert(pa->r() || (pa->d() == 1 && pa->p->r()));
        assert(!pa->is_vert);
        assert(pa->is_path);
        assert(pa->c[0] == this);
        assert(!pa->c[2]);

        if (pa->p) pa->p_c() = as_derived();
        this->p = pa->p;

        pa->is_path = false;
        pa->c[2] = pa->c[1];

        attach(pa, 0, c[0]);
        attach(pa, 1, c[1]);

        c[0] = 0;
        attach(as_derived(), 1, pa);
        assert(!c[2]);

        return pa->_pull(), pa;
    }

    ptr splice_non_path() {
        assert(!is_vert);
        assert(!is_path);

        splay();
        assert(p && p->is_vert && p->is_path);
        p->cut_right();

        if (!p->is_path) rot();
        assert(p && p->is_vert && p->is_path);
        assert(p->r() || (p->d() == 1 && p->p->r()));
        assert(p->c[d()] == this && !p->c[!d()]);

        ptr pa = p;

        if (pa->p) pa->p_c() = as_derived();
        this->p = pa->p;

        attach(pa, 0, c[0]);
        attach(pa, 1, c[1]);

        assert(c[2] && c[2]->is_path);
        attach(as_derived(), 0, pa);
        c[1] = c[2], c[2] = 0;

        is_path = true;

        return pa->_pull(), pa;
    }

    ptr splice_all() {
        ptr res = as_derived();
        for (ptr cur = as_derived(); cur; cur = cur->p) {
            if (!cur->is_path) res = cur->splice_non_path();
            assert(cur->is_path);
        }
        return res;
    }

  public:
    ptr expose() {
        assert(is_vert);

        push_all();

        ptr res = splice_all();
        cut_right(), pull_all();

        return res;
    }

    ptr expose_edge() {
        assert(!is_vert);

        push_all();

        ptr v = is_path ? c[1] : c[2];
        for (v->_push(); !v->is_vert; v->_push()) v = v->c[0];

        ptr res = v->splice_all();
        v->cut_right(), v->pull_all();

        assert(!p);
        assert(v == c[1]);

        return res == v ? as_derived() : res;
    }

    ptr meld_path_end() {
        assert(!p);
        ptr rt = as_derived();
        while (true) {
            rt->_push();
            if (rt->is_vert) break;
            rt = rt->c[1];
        }
        rt->splay_vert();
        if (rt->c[0] && rt->c[1]) {
            ptr ch = rt->c[1];
            while (true) {
                ch->_push();
                if (!ch->c[0]) break;
                ch = ch->c[0];
            }
            ch->splay();

            attach(ch, 0, rt->c[0]), rt->c[0] = 0;

            ch->_pull();
        } else if (rt->c[0]) {
            rt->c[1] = rt->c[0], rt->c[0] = 0;
        }
        return rt->pull_all();
    }

    void evert() {
        expose();

        ptr rt = as_derived();
        while (rt->p) {
            assert(rt->d() == 1);
            rt = rt->p;
        }

        rt->_flip();
        rt->meld_path_end();

        expose();

        assert(!p);
    }

    friend void link(ptr e, ptr v, ptr p) {
        assert(e && v && p);
        assert(!e->c[0] && !e->c[1] && !e->c[2]);
        v->evert(), p->expose();
        while (p->p) p = p->p;

        assert(!v->p);

        e->is_path = true, e->is_vert = false;
        attach(e, 0, p);
        attach(e, 1, v);
        e->_pull();
    }
    friend void link(node &e, node &v, node &p) { link(&e, &v, &p); }

    // Link v's root as a child of p with edge e
    // Returns false if they're already in the same subtree
    friend bool link_root(ptr e, ptr v, ptr p) {
        assert(e && v && p);
        assert(!e->c[0] && !e->c[1] && !e->c[2]);
        v->expose();
        p->expose();

        while (v->p) v = v->p;
        while (p->p) p = p->p;
        if (v == p) return false;

        e->is_path = true, e->is_vert = false;
        attach(e, 0, p);
        attach(e, 1, v);
        e->_pull();

        return true;
    }
    friend bool link_root(node &e, node &v, node &p) {
        return link_root(&e, &v, &p);
    }

    friend pair<ptr, ptr> cut(ptr e) {
        assert(!e->is_vert);
        e->expose_edge();

        assert(!e->p);
        assert(e->is_path);

        ptr l = e->c[0], r = e->c[1];
        assert(l && r);
        assert(!e->c[2]);

        e->c[0] = e->c[1] = 0;
        l->p = r->p = 0;

        l->meld_path_end();

        return {l, r};
    }
    friend pair<node &, node &> cut(node &e) {
        auto r = cut(&e);
        return {*r.first, *r.second};
    }

    friend ptr get_path(ptr v) {
        assert(v->is_vert);
        v->expose();
        if (!v->p) return v;
        assert(!v->p->p);
        return v->p;
    }
    friend node &get_path(node &v) { return *get_path(&v); }

    friend ptr get_subtree(ptr v) { return v->expose(), v; }
    friend node &get_subtree(node &v) { return v.expose(), v; }

    friend ptr lca(ptr u, ptr v) {
        if (u == v) return u;
        u->expose();
        auto up = u->p;
        ptr l = v->expose();
        if (u != v && up == u->p && (!up || !up->p)) return 0;
        return l;
    }
    friend ptr lca(node &u, node &v) { return lca(&u, &v); }

    template<class F> friend node *search(node *p, F &&select) {
        while (node *nxt = select(p)) {
            assert(p->is_vert);
            while (true) {
                p->_push(), nxt = select(p);
                if (nxt == p->c[2]) break;
                p = nxt;
            }
            p = nxt;
            assert(p->is_path);
            while (!p->is_vert) p->_push(), p = select(p);
        }
        p->evert();
        assert(!select(p));
        return p;
    }
    template<class F> friend node *search(node &p, F &&select) {
        return search(&p, forward<F>(select));
    }
};
#line 8 "test/yosupo/tree/point_set_tree_path_composite_sum_top_tree_ecnerwala.test.cpp"

using mint = modint998;
struct affine {
    mint a = 1, b = 0;
    affine operator()(const affine &g) const { return {a * g.a, a * g.b + b}; }
    mint operator()(mint x) const { return a * x + b; }
};

struct node : top_tree_node_base<node> {
    mint val;
    affine f;
    int subtree_size;
    array<affine, 2> path_affine;
    array<mint, 2> ans;

    void flip() { swap(path_affine[0], path_affine[1]), swap(ans[0], ans[1]); }
    void pull() {
        subtree_size = is_vert;
        path_affine.fill({});
        ans.fill(0);
        for (int i = 0; i <= 2; i++) {
            if (!c[i]) continue;
            subtree_size += c[i]->subtree_size;
        }
        if (is_vert) { // add_vertex(self, rake(c[0], c[1]))
            ans[0] = val;
            for (int i = 0; i < 2; i++) {
                if (!c[i]) continue;
                ans[0] += c[i]->ans[0];
            }
            ans[1] = ans[0];
        } else if (is_path) { // compress(c[0], self, c[1])
            for (int i = 0; i < 2; i++) {
                affine tmp = c[i]->path_affine[i](f);
                path_affine[i] = tmp(c[!i]->path_affine[i]);
                ans[i] = c[i]->ans[i] + tmp.a * c[!i]->ans[i] +
                         tmp.b * c[!i]->subtree_size;
            }
        } else { // rake(c[0], add_edge(self, c[2]), c[1])
            for (int i = 0; i < 2; i++) {
                if (!c[i]) continue;
                ans[0] += c[i]->ans[0];
            }
            ans[0] += f.a * c[2]->ans[0] + f.b * c[2]->subtree_size;
            ans[1] = ans[0];
        }
    }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<node> t(2 * n - 1);
    for (int i = 0; i < n; i++) {
        cin >> t[i].val;
        t[i].set_vert();
    }
    for (int i = 0, u, v; i < n - 1; i++) {
        cin >> u >> v >> t[i + n].f.a >> t[i + n].f.b;
        link(t[i + n], t[u], t[v]);
    }
    for (int qi = 0; qi < q; qi++) {
        int op, r;
        cin >> op;
        if (op == 0) {
            int w, x;
            cin >> w >> x;
            t[w].expose();
            t[w].val = x;
            t[w].pull_all();
        } else if (op == 1) {
            int e, y, z;
            cin >> e >> y >> z;
            auto &edge = t[e + n];
            edge.expose_edge();
            edge.f = {y, z};
            edge.pull_all();
        }
        cin >> r;
        t[r].evert();
        cout << t[r].ans[0] << "\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	3 MB
g++	typical_tree_max_00	AC	1079 ms	38 MB
g++	typical_tree_max_01	AC	659 ms	40 MB
g++	typical_tree_max_02	AC	1199 ms	38 MB
g++	typical_tree_max_03	AC	1086 ms	38 MB
g++	typical_tree_max_04	AC	1023 ms	38 MB
g++	typical_tree_max_05	AC	807 ms	38 MB
g++	typical_tree_max_06	AC	875 ms	38 MB
g++	typical_tree_max_07	AC	1018 ms	38 MB
g++	typical_tree_max_08	AC	984 ms	38 MB
g++	typical_tree_max_09	AC	722 ms	38 MB
g++	typical_tree_max_degree_00	AC	1009 ms	38 MB
g++	typical_tree_max_degree_01	AC	499 ms	39 MB
g++	typical_tree_max_degree_02	AC	1094 ms	38 MB
g++	typical_tree_max_degree_03	AC	1070 ms	38 MB
g++	typical_tree_max_degree_04	AC	1000 ms	38 MB
g++	typical_tree_max_degree_05	AC	724 ms	38 MB
g++	typical_tree_max_degree_06	AC	827 ms	38 MB
g++	typical_tree_max_degree_07	AC	1032 ms	38 MB
g++	typical_tree_max_degree_08	AC	983 ms	38 MB
g++	typical_tree_max_degree_09	AC	651 ms	38 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	990 ms	38 MB
clang++	typical_tree_max_01	AC	606 ms	41 MB
clang++	typical_tree_max_02	AC	1175 ms	38 MB
clang++	typical_tree_max_03	AC	1092 ms	38 MB
clang++	typical_tree_max_04	AC	975 ms	38 MB
clang++	typical_tree_max_05	AC	760 ms	38 MB
clang++	typical_tree_max_06	AC	863 ms	39 MB
clang++	typical_tree_max_07	AC	1024 ms	38 MB
clang++	typical_tree_max_08	AC	999 ms	38 MB
clang++	typical_tree_max_09	AC	670 ms	38 MB
clang++	typical_tree_max_degree_00	AC	969 ms	38 MB
clang++	typical_tree_max_degree_01	AC	464 ms	39 MB
clang++	typical_tree_max_degree_02	AC	1105 ms	38 MB
clang++	typical_tree_max_degree_03	AC	1093 ms	38 MB
clang++	typical_tree_max_degree_04	AC	974 ms	38 MB
clang++	typical_tree_max_degree_05	AC	685 ms	38 MB
clang++	typical_tree_max_degree_06	AC	793 ms	39 MB
clang++	typical_tree_max_degree_07	AC	1006 ms	38 MB
clang++	typical_tree_max_degree_08	AC	985 ms	38 MB
clang++	typical_tree_max_degree_09	AC	612 ms	38 MB

Back to top page