sotanishy's code snippets for competitive programming
View the Project on GitHub sotanishy/cp-library-cpp
#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"; } } }