test/heavy_light_decomposition_vertex_add_subtree_sum.test.cpp

• View this file on GitHub
• Last update: 2022-09-30 17:14:09+00:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

• Heavy-Light Decomposition (heavy_light_decomposition/heavy_light_decomposition.hpp)

Code

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

#include "heavy_light_decomposition/heavy_light_decomposition.hpp"
#include <atcoder/fenwicktree>
#include <bits/stdc++.h>

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    int N, Q;
    std::cin >> N >> Q;
    std::vector<int> a(N);
    for (auto &x : a) {
        std::cin >> x;
    }
    std::vector<std::vector<int>> adj(N);
    for (auto i = 1; i < N; ++i) {
        int p;
        std::cin >> p;
        adj[p].push_back(i);
    }
    heavy_light_decomposition hld(adj, 0);
    atcoder::fenwick_tree<long long> ft(N);
    for (auto i = 0; i < N; ++i) {
        hld.access_node(i, [&](auto x) { ft.add(x, a[i]); });
    }
    for (auto i = 0; i < Q; ++i) {
        int t, u;
        std::cin >> t >> u;
        if (t == 0) {
            long long val;
            std::cin >> val;
            hld.access_node(u, [&](auto x) { ft.add(x, val); });
        } else {
            auto ans = 0LL;
            hld.access_subtree(u, true,
                               [&](auto l, auto r) { ans = ft.sum(l, r); });
            std::cout << ans << '\n';
        }
    }
}

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

#line 1 "heavy_light_decomposition/heavy_light_decomposition.hpp"



#include <algorithm>
#include <cassert>
#include <type_traits>
#include <vector>

struct heavy_light_decomposition {
    heavy_light_decomposition(const std::vector<std::vector<int>> &graph,
                              int root)
        : n_(static_cast<int>(graph.size())), timer_(0), graph_(graph),
          size_(n_, 1), depth_(n_), parent_(n_, -1), top_(n_), in_(n_),
          out_(n_) {
        assert(0 <= root && root < n_);
        top_[root] = root;
        dfs_size(root);
        dfs_hld(root);
    }
    template <typename Function> void access_node(int u, Function f) {
        assert(0 <= u && u < n_);
        f(in_[u]);
    }
    template <typename Function>
    std::enable_if_t<
        std::is_same_v<std::invoke_result_t<Function, int, int>, void>, void>
    access_path(int u, int v, bool include_lca, Function f) const {
        assert(0 <= u && u < n_ && 0 <= v && v < n_);
        while (top_[u] != top_[v]) {
            if (depth_[top_[u]] < depth_[top_[v]]) {
                std::swap(u, v);
            }
            f(in_[top_[u]], in_[u] + 1);
            u = parent_[top_[u]];
        }
        if (depth_[u] > depth_[v]) {
            std::swap(u, v);
        }
        if (include_lca) {
            f(in_[u], in_[v] + 1);
        } else {
            f(in_[u] + 1, in_[v] + 1);
        }
    }
    template <typename Function>
    std::enable_if_t<
        std::is_same_v<std::invoke_result_t<Function, int, int, bool>, void>,
        void>
    access_path(int u, int v, bool include_lca, Function f) const {
        assert(0 <= u && u < n_ && 0 <= v && v < n_);
        bool u_to_lca = true;
        while (top_[u] != top_[v]) {
            if (depth_[top_[u]] < depth_[top_[v]]) {
                std::swap(u, v);
                u_to_lca = !u_to_lca;
            }
            f(in_[top_[u]], in_[u] + 1, u_to_lca);
            u = parent_[top_[u]];
        }
        if (depth_[u] > depth_[v]) {
            std::swap(u, v);
        } else {
            u_to_lca = !u_to_lca;
        }
        if (include_lca) {
            f(in_[u], in_[v] + 1, u_to_lca);
        } else {
            f(in_[u] + 1, in_[v] + 1, u_to_lca);
        }
    }
    template <typename Function>
    void access_subtree(int u, bool include_root, Function f) const {
        assert(0 <= u && u < n_);
        if (include_root) {
            f(in_[u], out_[u]);
        } else {
            f(in_[u] + 1, out_[u]);
        }
    }
    int lca(int u, int v) const {
        assert(0 <= u && u < n_ && 0 <= v && v < n_);
        while (top_[u] != top_[v]) {
            if (depth_[top_[u]] < depth_[top_[v]]) {
                std::swap(u, v);
            }
            u = parent_[top_[u]];
        }
        if (depth_[u] > depth_[v]) {
            std::swap(u, v);
        }
        return u;
    }

private:
    void dfs_size(int u) {
        auto it = std::find(graph_[u].begin(), graph_[u].end(), parent_[u]);
        if (it != graph_[u].end()) {
            graph_[u].erase(it);
        }
        for (auto &v : graph_[u]) {
            depth_[v] = depth_[u] + 1;
            parent_[v] = u;
            dfs_size(v);
            size_[u] += size_[v];
            if (size_[v] > size_[graph_[u][0]]) {
                std::swap(v, graph_[u][0]);
            }
        }
    }
    void dfs_hld(int u) {
        in_[u] = timer_++;
        for (auto v : graph_[u]) {
            top_[v] = (v == graph_[u][0] ? top_[u] : v);
            dfs_hld(v);
        }
        out_[u] = timer_;
    }
    int n_, timer_;
    std::vector<std::vector<int>> graph_;
    std::vector<int> size_, depth_, parent_, top_, in_, out_;
};


#line 4 "test/heavy_light_decomposition_vertex_add_subtree_sum.test.cpp"
#include <atcoder/fenwicktree>
#include <bits/stdc++.h>

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    int N, Q;
    std::cin >> N >> Q;
    std::vector<int> a(N);
    for (auto &x : a) {
        std::cin >> x;
    }
    std::vector<std::vector<int>> adj(N);
    for (auto i = 1; i < N; ++i) {
        int p;
        std::cin >> p;
        adj[p].push_back(i);
    }
    heavy_light_decomposition hld(adj, 0);
    atcoder::fenwick_tree<long long> ft(N);
    for (auto i = 0; i < N; ++i) {
        hld.access_node(i, [&](auto x) { ft.add(x, a[i]); });
    }
    for (auto i = 0; i < Q; ++i) {
        int t, u;
        std::cin >> t >> u;
        if (t == 0) {
            long long val;
            std::cin >> val;
            hld.access_node(u, [&](auto x) { ft.add(x, val); });
        } else {
            auto ans = 0LL;
            hld.access_subtree(u, true,
                               [&](auto l, auto r) { ans = ft.sum(l, r); });
            std::cout << ans << '\n';
        }
    }
}

Back to top page