test/yosupo/tree/dynamic_tree_subtree_add_subtree_sum_ett.test.cpp

View this file on GitHub
Last update: 2025-09-25 11:45:13+07:00
Problem: https://judge.yosupo.jp/problem/dynamic_tree_subtree_add_subtree_sum

Depends on

Code

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

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

#include "tree/lazy_ett.hpp"
#include "other/types.hpp"

struct M {
    struct T {
        int64_t sm = 0;
        int sz = 0;
    };
    using F = int64_t;
    static T id() { return {}; }
    static F fid() { return 0; }
    static T op(const T &l, const T &r) { return {l.sm + r.sm, l.sz + r.sz}; }
    static F comp(F l, F r) { return l + r; }
    static T map(F f, const T &x) { return {x.sm + f * x.sz, x.sz}; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    i32 n, q;
    cin >> n >> q;
    LazyETT<M> tr(n);
    for (i32 i = 0, a; i < n; i++) cin >> a, tr.set(i, {a, 1});
    for (i32 i = 0, u, v; i < n - 1; i++) cin >> u >> v, tr.link(u, v);
    while (q--) {
        i32 t;
        cin >> t;
        if (t == 0) {
            i32 a, b, c, d;
            cin >> a >> b >> c >> d;
            tr.cut(a, b), tr.link(c, d);
        } else if (t == 1) {
            i32 v, p, x;
            cin >> v >> p >> x;
            tr.apply(v, p, x);
        } else {
            i32 v, p;
            cin >> v >> p;
            cout << tr.subtree_prod(v, p).sm << "\n";
        }
    }
}

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

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

#line 2 "tree/lazy_ett.hpp"

#line 2 "other/st_alloc.hpp"

template<class T> struct st_alloc {
    static inline stack<T> pool;
    void *operator new(size_t) { return pool.emplace(), &pool.top(); }
    void operator delete(void *) {}
};
#line 4 "tree/lazy_ett.hpp"

template<class M> struct LazyETT {
    using T = typename M::T;
    using F = typename M::F;

    struct node;
    using ptr = node *;
    struct node : st_alloc<node> {
        ptr p = 0, l = 0, r = 0;
        int sz = 1;
        uint64_t pr = gen_pr();
        T val = M::id(), sm = M::id();
        bool upd = false;
        F lz = M::fid();

        node() {}

        static uint64_t gen_pr() {
            static mt19937_64 mt((size_t)make_unique<char>().get());
            return mt();
        }
    };
    static int sz(ptr t) { return t ? t->sz : 0; }
    static ptr &par(ptr t) {
        static ptr null = new node();
        return t ? t->p : null;
    }
    static ptr pull(ptr t) {
        t->sz = sz(t->l) + 1 + sz(t->r), t->sm = t->val;
        if (t->l) t->sm = M::op(t->l->sm, t->sm);
        if (t->r) t->sm = M::op(t->sm, t->r->sm);
        return t;
    }
    static void all_apply(ptr t, const F &f) {
        if (!t) return;
        t->val = M::map(f, t->val), t->sm = M::map(f, t->sm);
        t->lz = M::comp(f, t->lz), t->upd = true;
    }
    static void push(ptr t) {
        if (!t->upd) return;
        all_apply(t->l, t->lz), all_apply(t->r, t->lz);
        pull(t);
        t->lz = M::fid(), t->upd = false;
    }
    static int ord(ptr t) {
        int res = sz(t->l);
        for (; t->p; t = t->p)
            if (t == t->p->r) res += sz(t->p->l) + 1;
        return res;
    }
    static ptr root(ptr t) {
        while (t->p) t = t->p;
        return t;
    }
    static ptr merge(ptr l, ptr r) {
        if (!l || !r) return l ? l : r;
        push(l), push(r);
        if (l->pr < r->pr) {
            l->r = merge(l->r, r);
            par(l->r) = l;
            return pull(l);
        } else {
            r->l = merge(l, r->l);
            par(r->l) = r;
            return pull(r);
        }
    }
    static pair<ptr, ptr> split(ptr t, int k) {
        if (!t) return {};
        push(t);
        int w = sz(t->l) + 1;
        ptr a;
        if (k < w) {
            tie(a, t->l) = split(t->l, k);
            par(t->l) = t, par(a) = 0;
            return {a, pull(t)};
        } else {
            tie(t->r, a) = split(t->r, k - w);
            par(t->r) = t, par(a) = 0;
            return {pull(t), a};
        }
    }

    vector<map<int, ptr>> edges;
    ptr edge(int u, int v) { return edges[u][v] = edges[u][v] ?: new node(); }
    ptr vert(int v) { return edge(v, v); }
    ptr reroot(ptr e, bool rev) {
        auto [a, b] = split(root(e), ord(e) + rev);
        return merge(b, a);
    }

    LazyETT(int n) : edges(n) {}
    LazyETT(const vector<T> &val) : LazyETT((int)val.size()) {
        for (int i = 0; i < val.size(); i++)
            vert(i)->val = vert(i)->sm = val[i];
    }

    bool same(int u, int v) { return root(vert(u)) == root(vert(v)); }
    void link(int u, int v) {
        assert(!same(u, v));
        ptr a = reroot(vert(u), 0), b = reroot(vert(v), 1);
        merge(merge(a, edge(u, v)), merge(b, edge(v, u)));
    }
    void cut(int u, int v) {
        assert(same(u, v));
        ptr a, b;
        tie(a, b) = split(reroot(edge(u, v), 0), 1);
        tie(a, b) = split(b, ord(edge(v, u)));
        split(b, 1);
    }

    void set(int v, const T &x) {
        ptr p = vert(v);
        stack<ptr> st;
        while (p->p) p = p->p, st.push(p);
        while (st.size()) push(st.top()), st.pop();
        p = vert(v), p->val = x;
        for (pull(p); p->p;) pull(p = p->p);
    }
    T subtree_prod(int v, int p) {
        if (p != -1) cut(p, v);
        T ret = root(vert(v))->sm;
        if (p != -1) link(p, v);
        return ret;
    }
    void apply(int v, int p, const F &f) {
        if (p != -1) cut(p, v);
        all_apply(root(vert(v)), f);
        if (p != -1) link(p, v);
    }
};
#line 2 "other/types.hpp"

using i8 = int8_t;
using u8 = uint8_t;
using i32 = int32_t;
using u32 = uint32_t;
using i64 = int64_t;
using u64 = uint64_t;
using f32 = float;
using f64 = double;
using f80 = long double;
using str = string;
template<class T> using vec = vector<T>;
template<class E = i32> using graph = vec<vec<E>>;
#line 8 "test/yosupo/tree/dynamic_tree_subtree_add_subtree_sum_ett.test.cpp"

struct M {
    struct T {
        int64_t sm = 0;
        int sz = 0;
    };
    using F = int64_t;
    static T id() { return {}; }
    static F fid() { return 0; }
    static T op(const T &l, const T &r) { return {l.sm + r.sm, l.sz + r.sz}; }
    static F comp(F l, F r) { return l + r; }
    static T map(F f, const T &x) { return {x.sm + f * x.sz, x.sz}; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    i32 n, q;
    cin >> n >> q;
    LazyETT<M> tr(n);
    for (i32 i = 0, a; i < n; i++) cin >> a, tr.set(i, {a, 1});
    for (i32 i = 0, u, v; i < n - 1; i++) cin >> u >> v, tr.link(u, v);
    while (q--) {
        i32 t;
        cin >> t;
        if (t == 0) {
            i32 a, b, c, d;
            cin >> a >> b >> c >> d;
            tr.cut(a, b), tr.link(c, d);
        } else if (t == 1) {
            i32 v, p, x;
            cin >> v >> p >> x;
            tr.apply(v, p, x);
        } else {
            i32 v, p;
            cin >> v >> p;
            cout << tr.subtree_prod(v, p).sm << "\n";
        }
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	almost_line_00	AC	1643 ms	122 MB
g++	almost_line_01	AC	1599 ms	122 MB
g++	example_00	AC	5 ms	4 MB
g++	max_random_00	AC	1669 ms	124 MB
g++	max_random_01	AC	1562 ms	124 MB
g++	max_random_02	AC	1568 ms	124 MB
g++	max_random_03	AC	1553 ms	124 MB
g++	max_random_04	AC	1558 ms	124 MB
g++	random_00	AC	1104 ms	83 MB
g++	random_01	AC	1102 ms	94 MB
g++	random_02	AC	642 ms	42 MB
g++	random_03	AC	495 ms	90 MB
g++	random_04	AC	380 ms	21 MB
g++	small_00	AC	6 ms	4 MB
g++	small_01	AC	5 ms	4 MB
g++	small_02	AC	5 ms	4 MB
g++	small_03	AC	5 ms	4 MB
g++	small_04	AC	6 ms	4 MB
clang++	almost_line_00	AC	1705 ms	122 MB
clang++	almost_line_01	AC	1700 ms	122 MB
clang++	example_00	AC	5 ms	4 MB
clang++	max_random_00	AC	1701 ms	124 MB
clang++	max_random_01	AC	1667 ms	124 MB
clang++	max_random_02	AC	1620 ms	124 MB
clang++	max_random_03	AC	1636 ms	124 MB
clang++	max_random_04	AC	1597 ms	124 MB
clang++	random_00	AC	1065 ms	83 MB
clang++	random_01	AC	1125 ms	94 MB
clang++	random_02	AC	567 ms	42 MB
clang++	random_03	AC	489 ms	90 MB
clang++	random_04	AC	351 ms	21 MB
clang++	small_00	AC	6 ms	4 MB
clang++	small_01	AC	5 ms	4 MB
clang++	small_02	AC	5 ms	4 MB
clang++	small_03	AC	5 ms	4 MB
clang++	small_04	AC	6 ms	4 MB