test/yosupo/vertex_add_range_contour_sum_on_tree.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_range_contour_sum_on_tree

Depends on

• Fenwick Tree (data-structure/fenwick_tree.hpp)
• Centroid Decomposition (tree/centroid_decomposition.hpp)
• Range Contour Aggregation (tree/range_contour_aggregation.hpp)

Code

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

#include <bits/stdc++.h>

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

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

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

    int N, Q;
    cin >> N >> Q;
    std::vector<ll> a(N);
    for (auto& x : a) cin >> x;
    vector<vector<int>> G(N);
    for (int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }
    RangeContourAggregation<AddGroup> agg(G);
    for (int i = 0; i < N; ++i) {
        agg.update(i, a[i]);
    }
    for (int i = 0; i < Q; ++i) {
        int t, p;
        cin >> t >> p;
        if (t == 0) {
            ll x;
            cin >> x;
            agg.update(p, x);
        } else {
            int l, r;
            cin >> l >> r;
            ll ans = agg.fold(p, r) - agg.fold(p, l);
            cout << ans << "\n";
        }
    }
}

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

#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/centroid_decomposition.hpp"

std::tuple<std::vector<int>, std::vector<int>, std::vector<int>>
centroid_decomposition(const std::vector<std::vector<int>>& G) {
    const int N = G.size();
    std::vector<int> sz(N), level(N, -1), sz_comp(N), par(N);

    auto dfs_size = [&](auto& dfs_size, int v, int p) -> int {
        sz[v] = 1;
        for (int c : G[v]) {
            if (c != p && level[c] == -1) sz[v] += dfs_size(dfs_size, c, v);
        }
        return sz[v];
    };

    auto dfs_centroid = [&](auto& dfs_centroid, int v, int p, int n) -> int {
        for (int c : G[v]) {
            if (c != p && level[c] == -1 && sz[c] > n / 2)
                return dfs_centroid(dfs_centroid, c, v, n);
        }
        return v;
    };

    auto decompose = [&](auto& decompose, int v, int k, int p) -> void {
        int n = dfs_size(dfs_size, v, -1);
        int s = dfs_centroid(dfs_centroid, v, -1, n);
        level[s] = k;
        sz_comp[s] = n;
        par[s] = p;
        for (int c : G[s]) {
            if (level[c] == -1) decompose(decompose, c, k + 1, s);
        }
    };

    decompose(decompose, 0, 0, -1);
    return {level, sz_comp, par};
}
#line 6 "tree/range_contour_aggregation.hpp"

#line 9 "tree/range_contour_aggregation.hpp"

template <typename Group>
class RangeContourAggregation {
    using T = Group::T;

   public:
    RangeContourAggregation() = default;
    explicit RangeContourAggregation(const std::vector<std::vector<int>>& G)
        : seg(G.size()),
          seg_sub(G.size()),
          dist(G.size()),
          idx(G.size()),
          sub(G.size()),
          idx_sub(G.size()),
          first(G.size()),
          first_sub(G.size()) {
        std::tie(level, sz_comp, par) = centroid_decomposition(G);

        for (int v = 0; v < (int)G.size(); ++v) {
            // calculate for each vertex u,
            // dist: dist from v
            // idx: index in the BFS order
            // sub: subtree of v that u belongs to
            // idx_sub: index in the subtree

            // calculate for each depth d,
            // first: first index of depth d
            // first_sub: first index of depth d in the subtree i

            dist[v][v] = 0;
            std::queue<std::pair<int, int>> que;
            que.push({v, -1});
            int k = 0, sub_idx = 0;
            std::vector<int> ks;

            while (!que.empty()) {
                auto [u, i] = que.front();
                que.pop();

                // update component info
                int d = dist[v][u];
                if (d >= (int)first[v].size()) {
                    first[v].push_back(k);
                }
                idx[v][u] = k++;
                // enter subtree
                if (d == 1) {
                    i = sub_idx++;
                    ks.push_back(0);
                    first_sub[v].emplace_back();
                }
                // update subtree info
                if (i != -1) {
                    sub[v][u] = i;
                    if (d - 1 >= (int)first_sub[v][i].size()) {
                        first_sub[v][i].push_back(ks[i]);
                    }
                    idx_sub[v][u] = ks[i]++;
                }

                for (int w : G[u]) {
                    if (level[w] > level[v] && !dist[v].contains(w)) {
                        dist[v][w] = d + 1;
                        que.push({w, i});
                    }
                }
            }

            first[v].push_back(k);
            seg[v] = FenwickTree<Group>(k);
            for (int i = 0; i < sub_idx; ++i) {
                first_sub[v][i].push_back(ks[i]);
                seg_sub[v].emplace_back(ks[i]);
            }
        }
    }

    void update(int v, T x) {
        seg[v].update(0, x);
        for (int p = par[v]; p != -1; p = par[p]) {
            seg[p].update(idx[p][v], x);
            int i = sub[p][v];
            seg_sub[p][i].update(idx_sub[p][v], x);
        }
    }

    T fold(int v, int d) const {
        int dd = std::min((int)first[v].size() - 1, d);
        T res = seg[v].prefix_fold(first[v][dd]);
        for (int p = par[v]; p != -1; p = par[p]) {
            int dep = d - dist[p].at(v);
            if (dep >= 0) {
                dd = std::min((int)first[p].size() - 1, dep);
                res = Group::op(res, seg[p].prefix_fold(first[p][dd]));
                if (dd > 0) {
                    int i = sub[p].at(v);
                    dd = std::min((int)first_sub[p][i].size() - 1, dep - 1);
                    res = Group::op(res, Group::inv(seg_sub[p][i].prefix_fold(
                                             first_sub[p][i][dd])));
                }
            }
        }
        return res;
    }

   private:
    std::vector<int> level, sz_comp, par;
    std::vector<FenwickTree<Group>> seg;
    std::vector<std::vector<FenwickTree<Group>>> seg_sub;
    std::vector<std::unordered_map<int, int>> dist, idx, sub, idx_sub;
    std::vector<std::vector<int>> first;
    std::vector<std::vector<std::vector<int>>> first_sub;
};
#line 9 "test/yosupo/vertex_add_range_contour_sum_on_tree.test.cpp"
using namespace std;
using ll = long long;

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

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

    int N, Q;
    cin >> N >> Q;
    std::vector<ll> a(N);
    for (auto& x : a) cin >> x;
    vector<vector<int>> G(N);
    for (int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }
    RangeContourAggregation<AddGroup> agg(G);
    for (int i = 0; i < N; ++i) {
        agg.update(i, a[i]);
    }
    for (int i = 0; i < Q; ++i) {
        int t, p;
        cin >> t >> p;
        if (t == 0) {
            ll x;
            cin >> x;
            agg.update(p, x);
        } else {
            int l, r;
            cin >> l >> r;
            ll ans = agg.fold(p, r) - agg.fold(p, l);
            cout << ans << "\n";
        }
    }
}

Back to top page