Test/Library Checker/Tree/vertex_add_subtree_sum01.test.cpp

• View this file on GitHub
• Last update: 2023-12-31 10:56:46+09:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

• 双対セグメント木 (DataStructure/segtree_dual.hpp)
• HL分解 (HLD, Heavy-Light-Decomposition) (Tree/hld.hpp)

Code

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

#include <bits/stdc++.h>

#include "Tree/hld.hpp"

#include "DataStructure/segtree_dual.hpp"


using namespace std;
using ll = long long;
ll op(ll af, ll bf){ return af + bf; }
ll e(){return 0;}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int N, Q, cmd, u, v;
    cin >> N >> Q;
    vector<ll> a(N), b(N);
    vector<int> par(N);
    vector<vector<int>> g(N);
    for(auto &&v : a) cin >> v;
    for(int i = 1; i < N; i++){
        cin >> par[i];
        g[par[i]].emplace_back(i);
    }
    Heavy_Light_Decomposition HLD(g);
    for(int i = N - 1; i >= 1; i--){
        a[par[i]] += a[i];
        b[HLD[i]] = a[i];
    }
    b[0] = a[0];
    dual_segtree<ll, op, e> seg(b);
    auto update = [&](int l, int r){
        seg.apply(l, r, v);
    };
    while(Q--){
        cin >> cmd >> u;
        if(cmd == 0){
            cin >> v;
            HLD.update(0, u, update);
        }else{
            cout << seg[HLD[u]] << '\n';
        }
    }
}

#line 1 "Test/Library Checker/Tree/vertex_add_subtree_sum01.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"

#include <bits/stdc++.h>

#line 1 "Tree/hld.hpp"
struct Heavy_Light_Decomposition{
    int N, tim = 0;
    std::vector<int> sz, ent, leader, order, par;
    std::vector<std::vector<int>> &G;
    Heavy_Light_Decomposition(std::vector<std::vector<int>> &g) : 
            N(g.size()), G(g), sz(N), ent(N), leader(N), order(N), par(N) {
        dfs_size(0, -1);
        dfs_hld(0);
    }
    const int operator[](int v) const {
        assert(0 <= v && v < N);
        return ent[v];
    }
    int operator[](int v) { 
        assert(0 <= v && v < N);
        return ent[v];
    }
    int la(int v, int k) {
        while(true) {
            int u = leader[v];
            if(ent[v] - k >= ent[u]) return order[ent[v] - k];
            k -= ent[v] - ent[u] + 1;
            v = par[u];
        }
    }
    int la(int from, int to, int d){
        int d1 = 0, d2 = 0;
        int v = from, u = to;
        do{
            if(ent[u] < ent[v]){
                if(leader[u] == leader[v]){
                    d1 += ent[v] - ent[u];
                    break;
                }
                d1 += ent[v] - ent[leader[v]] + 1;
                v = par[leader[v]];
            }else{
                if(leader[u] == leader[v]){
                    d2 += ent[u] - ent[v];
                    break;
                }
                d2 += ent[u] - ent[leader[u]] + 1;
                u = par[leader[u]];
            }
        }while(true);
        if(d > d1 + d2) return -1;
        return d <= d1 ? la(from, d) : la(to, d1 + d2 - d);
    }
    int lca(int u, int v) {
        do{
            if(ent[u] > ent[v]) std::swap(u, v);
            if(leader[u] == leader[v]) return u;
            v = par[leader[v]];
        }while(true);
    }
    int dist(int u, int v){
        int ans = 0;
        do{
            if(ent[u] > ent[v]) std::swap(u, v);
            if(leader[u] == leader[v]) return ans + ent[v] - ent[u];
            ans += ent[v] - ent[leader[v]] + 1;
            v = par[leader[v]];
        }while(true);
    }
    template< typename T, typename Q, typename F >
    T query(int u, int v, const T &identity, const Q &qf, const F &f, bool edge = false) {
        T ans = identity;
        do{
            if(ent[u] > ent[v]) std::swap(u, v);
            if(leader[u] == leader[v]) break;
            ans = f( qf(ent[leader[v]], ent[v] + 1), ans);
            v = par[leader[v]];
        }while(true);
        return f( qf(ent[u] + edge, ent[v] + 1), ans);
    }
    template< typename T, typename Q1, typename Q2, typename F >
    T noncom_query(int u, int v, const T &identity, 
                   const Q1 &qf, const Q2 &rev_qf, const F &f, bool edge = false) {
        T sml = identity, smr = identity;
        do{
            if(leader[u] == leader[v]) break;
            if(ent[u] < ent[v]){
                smr = f( qf(ent[leader[v]], ent[v] + 1), smr);
                v = par[leader[v]];
            }else{
                sml = f( sml, rev_qf(ent[leader[u]], ent[u] + 1));
                u = par[leader[u]];
            }
        }while(true);
        if(ent[u] < ent[v]){
            return f(sml,  f( qf(ent[u] + edge, ent[v] + 1), smr));
        }else{
            return f(f(sml, rev_qf(ent[v] + edge, ent[u] + 1)), smr);
        }
    }
    template< typename Q >
    void update(int u, int v, const Q &q, bool edge = false) {
        do{
            if(ent[u] > ent[v]) std::swap(u, v);
            if(leader[u] == leader[v]) break;
            q(ent[leader[v]], ent[v] + 1);
            v = par[leader[v]];
        }while(true);
        q(ent[u] + edge, ent[v] + 1);
    }
    private:
    void dfs_size(int v, int p){
        par[v] = p;
        sz[v] = 1;
        if(!G[v].empty() && G[v][0] == p) std::swap(G[v][0], G[v].back());
        for(auto &u : G[v]){
            if(u == p) continue;
            dfs_size(u, v);
            sz[v] += sz[u];
            if(sz[u] > sz[G[v][0]]) std::swap(G[v][0], u);
        }
    }
    void dfs_hld(int v){
        ent[v] = tim++;
        order[ent[v]] = v;
        for(auto &u : G[v]) {
            if(u == par[v]) continue;
            leader[u] = (G[v][0] == u ? leader[v] : u);
            dfs_hld(u);
        }
    }
};
#line 1 "DataStructure/segtree_dual.hpp"
template <class S, S (*mapping)(S, S), S (*id)()> struct dual_segtree {
    public:
    dual_segtree() : dual_segtree(0) {}
    dual_segtree(int n) : dual_segtree(std::vector<S>(n, id())) {}
    dual_segtree(const std::vector<S>& v) : _n(int(v.size())) {
        log = ceil_pow2(_n);
        size = 1 << log;
        d = std::vector<S>(2 * size, id());
        for (int i = 0; i < _n; i++) d[size + i] = v[i];
    }
    const S& operator[](int p) const {
        assert(0 <= p && p < _n);
        p += size;
        for (int i = log; i >= 1; i--) push(p >> i);
        return d[p];
    }
    S& operator[](int p) { 
        assert(0 <= p && p < _n);
        p += size;
        for (int i = log; i >= 1; i--) push(p >> i);
        return d[p];
    }
    void apply(int p, S f) {
        assert(0 <= p && p < _n);
        p += size;
        for (int i = log; i >= 1; i--) push(p >> i);
        d[p] = mapping(f, d[p]);
    }
    void apply(int l, int r, S f) {
        assert(0 <= l && l <= r && r <= _n);
        if (l == r) return;
        l += size;
        r += size;
        for (int i = log; i >= 1; i--) {
            if (((l >> i) << i) != l) push(l >> i);
            if (((r >> i) << i) != r) push((r - 1) >> i);
        }
        while (l < r) {
            if (l & 1) all_apply(l++, f);
            if (r & 1) all_apply(--r, f);
            l >>= 1;
            r >>= 1;
        }
    }
    private:
    int _n, size, log;
    std::vector<S> d;
    void all_apply(int k, S f) {
        d[k] = mapping(f, d[k]);
    }
    void push(int k) {
        all_apply(2 * k, d[k]);
        all_apply(2 * k + 1, d[k]);
        d[k] = id();
    }
    int ceil_pow2(int n) {
        int x = 0;
        while ((1U << x) < (unsigned int)(n)) x++;
        return x;
    }
};
#line 6 "Test/Library Checker/Tree/vertex_add_subtree_sum01.test.cpp"

using namespace std;
using ll = long long;
ll op(ll af, ll bf){ return af + bf; }
ll e(){return 0;}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int N, Q, cmd, u, v;
    cin >> N >> Q;
    vector<ll> a(N), b(N);
    vector<int> par(N);
    vector<vector<int>> g(N);
    for(auto &&v : a) cin >> v;
    for(int i = 1; i < N; i++){
        cin >> par[i];
        g[par[i]].emplace_back(i);
    }
    Heavy_Light_Decomposition HLD(g);
    for(int i = N - 1; i >= 1; i--){
        a[par[i]] += a[i];
        b[HLD[i]] = a[i];
    }
    b[0] = a[0];
    dual_segtree<ll, op, e> seg(b);
    auto update = [&](int l, int r){
        seg.apply(l, r, v);
    };
    while(Q--){
        cin >> cmd >> u;
        if(cmd == 0){
            cin >> v;
            HLD.update(0, u, update);
        }else{
            cout << seg[HLD[u]] << '\n';
        }
    }
}

Back to top page