tree/test/preorder_tree_dfs.test.cpp

• View this file on GitHub
• Last update: 2025-01-03 09:34:57+09:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

• Preorder tree DFS （木の行きがけ順 DFS） (tree/preorder_tree_dfs.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
#include "../preorder_tree_dfs.hpp"
#include <cassert>
#include <iostream>
#include <atcoder/fenwicktree>
using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<long long> A(N);
    vector<vector<int>> to(N);
    for (auto &a : A) cin >> a;
    for (int i = 1; i < N; i++) {
        int p;
        cin >> p;
        to.at(p).push_back(i);
        to.at(i).push_back(p);
    }

    preorder_tree_dfs tour(to, 0);

    atcoder::fenwick_tree<long long> ft(N);
    for (int i = 0; i < N; i++) ft.add(tour.subtree_begin.at(i), A.at(i));

    while (Q--) {
        int q;
        cin >> q;
        if (q) {
            int u;
            cin >> u;
            printf("%lld\n", ft.sum(tour.subtree_begin[u], tour.subtree_end[u]));
        } else {
            int u, x;
            cin >> u >> x;
            ft.add(tour.subtree_begin[u], x);
        }
    }
}

#line 1 "tree/test/preorder_tree_dfs.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
#line 2 "tree/preorder_tree_dfs.hpp"
#include <cassert>
#include <vector>

// Preorder tree DFS
// 木を DFS して行きがけ順に頂点を保持する．
// 木を区間に変換する．部分木を構成する頂点は連続して現れるので，部分木の頂点クエリ等に有用．
// heavy_light_decomposition が上位互換
struct preorder_tree_dfs {
    int n; // # of vertices of tree
    int root;
    std::vector<int> subtree_begin, subtree_end;
    std::vector<int> vis_order;

    void _build_dfs(int now, const std::vector<std::vector<int>> &to) {
        subtree_begin[now] = vis_order.size();
        vis_order.push_back(now);
        for (auto nxt : to[now]) {
            if (subtree_begin[nxt] == -1) _build_dfs(nxt, to);
        }
        subtree_end[now] = vis_order.size();
    }

    preorder_tree_dfs() = default;

    preorder_tree_dfs(const std::vector<std::vector<int>> &to, int root)
        : n(to.size()), root(root), subtree_begin(n, -1), subtree_end(n, -1) {
        _build_dfs(root, to);
    }
};
#line 4 "tree/test/preorder_tree_dfs.test.cpp"
#include <iostream>
#include <atcoder/fenwicktree>
using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<long long> A(N);
    vector<vector<int>> to(N);
    for (auto &a : A) cin >> a;
    for (int i = 1; i < N; i++) {
        int p;
        cin >> p;
        to.at(p).push_back(i);
        to.at(i).push_back(p);
    }

    preorder_tree_dfs tour(to, 0);

    atcoder::fenwick_tree<long long> ft(N);
    for (int i = 0; i < N; i++) ft.add(tour.subtree_begin.at(i), A.at(i));

    while (Q--) {
        int q;
        cin >> q;
        if (q) {
            int u;
            cin >> u;
            printf("%lld\n", ft.sum(tour.subtree_begin[u], tour.subtree_end[u]));
        } else {
            int u, x;
            cin >> u >> x;
            ft.add(tour.subtree_begin[u], x);
        }
    }
}

Back to top page