test/yosupo/vertex_add_subtree_sum.hld.test.cpp

View this file on GitHub
Last update: 2024-01-07 23:25:49+09:00
Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

Code

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

#include <bits/stdc++.h>

#include "../../data-structure/fenwick_tree.hpp"
#include "../../tree/hld.hpp"
using namespace std;
using ll = long long;

struct AddMonoid {
    using T = ll;
    static T id() { return 0; }
    static T op(T a, T b) { return a + b; }
};

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    vector<vector<int>> G(N);
    for (int i = 1; i < N; ++i) {
        int p;
        cin >> p;
        G[p].push_back(i);
    }
    HLD<AddMonoid> hld(G, false);
    FenwickTree<AddMonoid> fw(N);
    auto update = [&](int k, ll x) { fw.update(k, x); };
    auto fold = [&](int l, int r) {
        return fw.prefix_fold(r) - fw.prefix_fold(l);
    };
    for (int i = 0; i < N; ++i) hld.update(i, a[i], update);
    for (int i = 0; i < Q; ++i) {
        int t, u;
        cin >> t >> u;
        if (t == 0) {
            int x;
            cin >> x;
            hld.update(u, x, update);
        } else {
            cout << hld.subtree_fold(u, fold) << "\n";
        }
    }
}

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

#include <bits/stdc++.h>

#line 4 "data-structure/fenwick_tree.hpp"

template <typename M>
class FenwickTree {
    using T = M::T;

   public:
    FenwickTree() = default;
    explicit FenwickTree(int n) : n(n), data(n + 1, M::id()) {}

    T prefix_fold(int i) const {
        T ret = M::id();
        for (; i > 0; i -= i & -i) ret = M::op(ret, data[i]);
        return ret;
    }

    void update(int i, const T& x) {
        for (++i; i <= n; i += i & -i) data[i] = M::op(data[i], x);
    }

    int lower_bound(const T& x) const { return lower_bound(x, std::less<>()); }

    template <typename Compare>
    int lower_bound(const T& x, Compare cmp) const {
        if (!cmp(M::id(), x)) return 0;
        int k = 1;
        while (k * 2 <= n) k <<= 1;
        int i = 0;
        T v = M::id();
        for (; k > 0; k >>= 1) {
            if (i + k > n) continue;
            T nv = M::op(v, data[i + k]);
            if (cmp(nv, x)) {
                v = nv;
                i += k;
            }
        }
        return i + 1;
    }

   private:
    int n;
    std::vector<T> data;
};
#line 4 "tree/hld.hpp"

template <typename M>
class HLD {
    using T = M::T;

   public:
    HLD() = default;
    HLD(const std::vector<std::vector<int>>& G, bool edge)
        : G(G),
          size(G.size()),
          depth(G.size()),
          par(G.size(), -1),
          in(G.size()),
          out(G.size()),
          head(G.size()),
          heavy(G.size(), -1),
          edge(edge) {
        dfs(0);
        decompose(0, 0);
    }

    template <typename F>
    void update(int v, const T& x, const F& f) const {
        f(in[v], x);
    }

    template <typename F>
    void update_edge(int u, int v, const T& x, const F& f) const {
        if (in[u] > in[v]) std::swap(u, v);
        f(in[v], x);
    }

    template <typename E, typename F>
    void update(int u, int v, const E& x, const F& f) const {
        while (head[u] != head[v]) {
            if (in[head[u]] > in[head[v]]) std::swap(u, v);
            f(in[head[v]], in[v] + 1, x);
            v = par[head[v]];
        }
        if (in[u] > in[v]) std::swap(u, v);
        f(in[u] + edge, in[v] + 1, x);
    }

    template <typename F, typename Flip>
    T path_fold(int u, int v, const F& f, const Flip& flip) const {
        bool flipped = false;
        T resu = M::id(), resv = M::id();
        while (head[u] != head[v]) {
            if (in[head[u]] > in[head[v]]) {
                std::swap(u, v);
                std::swap(resu, resv);
                flipped ^= true;
            }
            T val = f(in[head[v]], in[v] + 1);
            resv = M::op(val, resv);
            v = par[head[v]];
        }
        if (in[u] > in[v]) {
            std::swap(u, v);
            std::swap(resu, resv);
            flipped ^= true;
        }
        T val = f(in[u] + edge, in[v] + 1);
        resv = M::op(val, resv);
        resv = M::op(flip(resu), resv);
        if (flipped) {
            resv = flip(resv);
        }
        return resv;
    }

    template <typename F>
    T path_fold(int u, int v, const F& f) const {
        return path_fold(u, v, f, [&](auto& v) { return v; });
    }

    template <typename F>
    T subtree_fold(int v, const F& f) const {
        return f(in[v] + edge, out[v]);
    }

    int lca(int u, int v) const {
        while (true) {
            if (in[u] > in[v]) std::swap(u, v);
            if (head[u] == head[v]) return u;
            v = par[head[v]];
        }
    }

    int dist(int u, int v) const {
        return depth[u] + depth[v] - 2 * depth[lca(u, v)];
    }

   private:
    std::vector<std::vector<int>> G;
    std::vector<int> size, depth, par, in, out, head, heavy;
    bool edge;
    int cur_pos = 0;

    void dfs(int v) {
        size[v] = 1;
        int max_size = 0;
        for (int c : G[v]) {
            if (c == par[v]) continue;
            par[c] = v;
            depth[c] = depth[v] + 1;
            dfs(c);
            size[v] += size[c];
            if (size[c] > max_size) {
                max_size = size[c];
                heavy[v] = c;
            }
        }
    }

    void decompose(int v, int h) {
        head[v] = h;
        in[v] = cur_pos++;
        if (heavy[v] != -1) decompose(heavy[v], h);
        for (int c : G[v]) {
            if (c != par[v] && c != heavy[v]) decompose(c, c);
        }
        out[v] = cur_pos;
    }
};
#line 7 "test/yosupo/vertex_add_subtree_sum.hld.test.cpp"
using namespace std;
using ll = long long;

struct AddMonoid {
    using T = ll;
    static T id() { return 0; }
    static T op(T a, T b) { return a + b; }
};

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    vector<vector<int>> G(N);
    for (int i = 1; i < N; ++i) {
        int p;
        cin >> p;
        G[p].push_back(i);
    }
    HLD<AddMonoid> hld(G, false);
    FenwickTree<AddMonoid> fw(N);
    auto update = [&](int k, ll x) { fw.update(k, x); };
    auto fold = [&](int l, int r) {
        return fw.prefix_fold(r) - fw.prefix_fold(l);
    };
    for (int i = 0; i < N; ++i) hld.update(i, a[i], update);
    for (int i = 0; i < Q; ++i) {
        int t, u;
        cin >> t >> u;
        if (t == 0) {
            int x;
            cin >> x;
            hld.update(u, x, update);
        } else {
            cout << hld.subtree_fold(u, fold) << "\n";
        }
    }
}