test/yosupo/tree/dynamic_tree_vertex_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_vertex_add_subtree_sum

Depends on

Code

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

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

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

struct M {
    using T = i64;
    static T id() { return 0; }
    static T op(T l, T r) { return l + r; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    i32 n, q;
    cin >> n >> q;
    ETT<M> tr(n);
    for (i32 i = 0, a; i < n; i++) cin >> a, tr.set(i, a);
    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 p, x;
            cin >> p >> x;
            tr.set(p, tr.get(p) + x);
        } else {
            i32 v, p;
            cin >> v >> p;
            cout << tr.subtree_prod(v, p) << "\n";
        }
    }
}

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

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

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

template<class M> struct ETT {
    using T = typename M::T;
    struct node;
    using ptr = node *;
    struct node : st_alloc<node> {
        ptr p = 0, l = 0, r = 0;
        int sz = 1;
        T prd = M::id(), val = M::id();
        uint64_t pr = gen();

        node() {}
        node(const T &x) : prd(x), val(x) {}

        static uint64_t gen() {
            static mt19937_64 mt(uint64_t(make_unique<char>().get()));
            return mt();
        }
    };

    static int sz(ptr t) { return t ? t->sz : 0; }
    static T prd(ptr t) { return t ? t->prd : 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->prd = M::op(M::op(prd(t->l), t->val), prd(t->r));
        return t;
    }
    static ptr merge(ptr l, ptr r) {
        if (!l || !r) return l ? l : 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 {};
        int w = sz(t->l);
        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 - 1);
            par(t->r) = t, par(a) = 0;
            return {pull(t), a};
        }
    }
    static ptr root(ptr t) { return t->p ? root(t->p) : t; }
    static int ord(ptr t) {
        int o = sz(t->l);
        for (; t->p; t = t->p)
            if (t == t->p->r) o += sz(t->p->l) + 1;
        return o;
    }

    vector<map<int, ptr>> edges;

    ptr &edge(int u, int v) { return edges[u][v] = edges[u][v] ?: new node(); }
    ptr rot(int u, int v, bool rev) {
        ptr e = edge(u, v), a, b;
        tie(a, b) = split(root(e), ord(e) + rev);
        return merge(b, a);
    }

    ETT(int n) : edges(n) {}

    T get(int v) { return edge(v, v)->val; }
    void set(int v, const T &x) {
        ptr e = edge(v, v);
        for (e->val = x; e; e = e->p) pull(e);
    }
    bool connected(int u, int v) {
        return root(edge(u, u)) == root(edge(v, v));
    }
    void link(int u, int v) {
        if (connected(u, v)) return;
        ptr a = rot(u, u, 0), b = rot(v, v, 1);
        merge(merge(a, edge(u, v)), merge(b, edge(v, u)));
    }
    void cut(int u, int v) {
        if (!connected(u, v)) return;
        ptr a, b;
        tie(a, b) = split(rot(u, v, 0), 1);
        tie(a, b) = split(b, ord(edge(v, u)));
        split(b, 1);
    }
    int subtree_size(int v, int p = -1) {
        if (p != -1) cut(v, p);
        int ret = root(edge(v, v))->sz;
        if (p != -1) link(v, p);
        return ret;
    }
    T subtree_prod(int v, int p = -1) {
        if (p != -1) cut(v, p);
        T ret = root(edge(v, v))->prd;
        if (p != -1) link(v, p);
        return ret;
    }
};
#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_vertex_add_subtree_sum_ett.test.cpp"

struct M {
    using T = i64;
    static T id() { return 0; }
    static T op(T l, T r) { return l + r; }
};

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    i32 n, q;
    cin >> n >> q;
    ETT<M> tr(n);
    for (i32 i = 0, a; i < n; i++) cin >> a, tr.set(i, a);
    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 p, x;
            cin >> p >> x;
            tr.set(p, tr.get(p) + x);
        } else {
            i32 v, p;
            cin >> v >> p;
            cout << tr.subtree_prod(v, p) << "\n";
        }
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	almost_line_00	AC	1213 ms	98 MB
g++	almost_line_01	AC	1140 ms	98 MB
g++	example_00	AC	5 ms	4 MB
g++	max_random_00	AC	1327 ms	100 MB
g++	max_random_01	AC	1297 ms	100 MB
g++	max_random_02	AC	1318 ms	100 MB
g++	max_random_03	AC	1289 ms	100 MB
g++	max_random_04	AC	1302 ms	100 MB
g++	random_00	AC	826 ms	67 MB
g++	random_01	AC	871 ms	76 MB
g++	random_02	AC	398 ms	34 MB
g++	random_03	AC	480 ms	73 MB
g++	random_04	AC	239 ms	17 MB
g++	small_00	AC	7 ms	4 MB
g++	small_01	AC	6 ms	4 MB
g++	small_02	AC	6 ms	4 MB
g++	small_03	AC	7 ms	4 MB
g++	small_04	AC	6 ms	4 MB
clang++	almost_line_00	AC	1066 ms	98 MB
clang++	almost_line_01	AC	1026 ms	98 MB
clang++	example_00	AC	5 ms	4 MB
clang++	max_random_00	AC	1190 ms	100 MB
clang++	max_random_01	AC	1211 ms	100 MB
clang++	max_random_02	AC	1248 ms	100 MB
clang++	max_random_03	AC	1307 ms	100 MB
clang++	max_random_04	AC	1260 ms	100 MB
clang++	random_00	AC	810 ms	67 MB
clang++	random_01	AC	893 ms	76 MB
clang++	random_02	AC	435 ms	34 MB
clang++	random_03	AC	506 ms	73 MB
clang++	random_04	AC	230 ms	17 MB
clang++	small_00	AC	7 ms	4 MB
clang++	small_01	AC	7 ms	4 MB
clang++	small_02	AC	7 ms	4 MB
clang++	small_03	AC	7 ms	4 MB
clang++	small_04	AC	6 ms	4 MB