verify/LibraryChecker/tree/VertexAddSubtreeSum.test.cpp

• View this file on GitHub
• Last update: 2024-06-26 20:58:27+09:00
• Problem: https://judge.yosupo.jp/problem/vertex_add_subtree_sum

Depends on

• Sparse Table (data-structure/others/SparseTable.hpp)
• Segment Tree (data-structure/segment-tree/SegmentTree.hpp)
• Euler Tour (graph/tree/EulerTour.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
#include <bits/stdc++.h>
using namespace std;
#include "../../../graph/tree/EulerTour.hpp"

using S = long long;
S op(S a, S b) {
  return a + b;
}

S e() {
  return 0;
}

S inv(S a) {
  return -a;
}

int main() {
  cin.tie(0)->sync_with_stdio(0);
  int n, q;
  cin >> n >> q;
  vector<vector<pair<int, S>>> g(n);
  vector<S> nodew(n);
  for (int i = 0; i < n; i++) {
    cin >> nodew[i];
  }
  for (int i = 1; i < n; i++) {
    int v;
    cin >> v;
    g[i].push_back({v, 0});
    g[v].push_back({i, 0});
  }

  EulerTour<S, op, e, inv, S, op, e, inv> t(g, nodew);

  for (int i = 0; i < q; i++) {
    int com;
    cin >> com;
    if (com == 0) {
      int u, x;
      cin >> u >> x;
      t.update_weight_node(u, t.get_node(u) + x);
    } else if (com == 1) {
      int u;
      cin >> u;
      cout << t.subtree_node(u) << "\n";
    }
  }
}

#line 1 "verify/LibraryChecker/tree/VertexAddSubtreeSum.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
#include <bits/stdc++.h>
using namespace std;
#line 1 "data-structure/others/SparseTable.hpp"
template <class T, auto f>
class SparseTable {
 private:
  vector<vector<T>> st;
  vector<int> lookup;

 public:
  SparseTable() {}

  SparseTable(const vector<T> &v) {
    const int n = (int)v.size();
    const int b = 32 - __builtin_clz(n);
    st.assign(b, vector<T>(n));
    for (int i = 0; i < v.size(); i++) {
      st[0][i] = v[i];
    }
    
    for (int i = 1; i < b; i++) {
      for (int j = 0; j + (1 << i) <= n; j++) {
        st[i][j] = f(st[i - 1][j], st[i - 1][j + (1 << (i - 1))]);
      }
    }
    lookup.resize(v.size() + 1);
    for (int i = 2; i < lookup.size(); i++) {
      lookup[i] = lookup[i >> 1] + 1;
    }
  }

  inline T query(int l, int r) const {
    int b = lookup[r - l];
    return f(st[b][l], st[b][r - (1 << b)]);
  }
};
#line 2 "data-structure/segment-tree/SegmentTree.hpp"
template <class S, auto op, auto e>
struct segtree {
 private:
  unsigned int seg_bit_ceil(unsigned int n) {
    unsigned int x = 1;
    while (x < (unsigned int)(n)) x *= 2;
    return x;
  }

 public:
  static_assert(std::is_convertible_v<decltype(op), std::function<S(S, S)>>,
                "op must work as S(S, S)");
  static_assert(std::is_convertible_v<decltype(e), std::function<S()>>,
                "e must work as S()");
  segtree() : segtree(0) {}
  explicit segtree(int n) : segtree(std::vector<S>(n, e())) {}
  explicit segtree(const std::vector<S>& v) : _n(int(v.size())) {
    size = (int)seg_bit_ceil((unsigned int)(_n));
    log = __builtin_ctz((unsigned int)size);
    d = std::vector<S>(2 * size, e());
    for (int i = 0; i < _n; i++) d[size + i] = v[i];
    for (int i = size - 1; i >= 1; i--) {
      update(i);
    }
  }

  void set(int p, S x) {
    assert(0 <= p && p < _n);
    p += size;
    d[p] = x;
    for (int i = 1; i <= log; i++) update(p >> i);
  }

  S get(int p) const {
    assert(0 <= p && p < _n);
    return d[p + size];
  }

  S prod(int l, int r) const {
    assert(0 <= l && l <= r && r <= _n);
    S sml = e(), smr = e();
    l += size;
    r += size;

    while (l < r) {
      if (l & 1) sml = op(sml, d[l++]);
      if (r & 1) smr = op(d[--r], smr);
      l >>= 1;
      r >>= 1;
    }
    return op(sml, smr);
  }

  S all_prod() const { return d[1]; }

  template <bool (*f)(S)>
  int max_right(int l) const {
    return max_right(l, [](S x) { return f(x); });
  }
  template <class F>
  int max_right(int l, F f) const {
    assert(0 <= l && l <= _n);
    assert(f(e()));
    if (l == _n) return _n;
    l += size;
    S sm = e();
    do {
      while (l % 2 == 0) l >>= 1;
      if (!f(op(sm, d[l]))) {
        while (l < size) {
          l = (2 * l);
          if (f(op(sm, d[l]))) {
            sm = op(sm, d[l]);
            l++;
          }
        }
        return l - size;
      }
      sm = op(sm, d[l]);
      l++;
    } while ((l & -l) != l);
    return _n;
  }

  template <bool (*f)(S)>
  int min_left(int r) const {
    return min_left(r, [](S x) { return f(x); });
  }
  template <class F>
  int min_left(int r, F f) const {
    assert(0 <= r && r <= _n);
    assert(f(e()));
    if (r == 0) return 0;
    r += size;
    S sm = e();
    do {
      r--;
      while (r > 1 && (r % 2)) r >>= 1;
      if (!f(op(d[r], sm))) {
        while (r < size) {
          r = (2 * r + 1);
          if (f(op(d[r], sm))) {
            sm = op(d[r], sm);
            r--;
          }
        }
        return r + 1 - size;
      }
      sm = op(d[r], sm);
    } while ((r & -r) != r);
    return 0;
  }

 private:
  int _n, size, log;
  std::vector<S> d;

  void update(int k) { d[k] = op(d[2 * k], d[2 * k + 1]); }
};
#line 4 "graph/tree/EulerTour.hpp"
struct SD {
  int depth;
  int ind;
};
auto f = [](const SD x, const SD y) { if (x.depth < y.depth) return x; else return y; };
template <class S, auto ops, auto es, auto invs, class T, auto opt, auto et, auto invt>
class EulerTour {
 private:
  int n;
  vector<int> start, finish, parent, depth, euler_tour;
  vector<SD> depth_v;
  vector<S> node_sum;
  vector<T> edge_sum;
  segtree<T, opt, et> seg_e, pseg_e;
  segtree<S, ops, es> seg_n, pseg_n;
  SparseTable<SD, f> sp_d;
  int time = 0;

  void rec(int cur, int p, int d, vector<vector<pair<int, S>>> &g) {
    start[cur] = time++;
    parent[cur] = p;
    depth[cur] = d;
    for (auto nex : g[cur]) {
      if (nex.first != p) {
        rec(nex.first, cur, d + 1, g);
      }
    }
    finish[cur] = time++;
  }

  T root_path_edge(int ind) {
    return pseg_e.prod(1, start[ind] + 1);
  }

  S root_path_node(int ind) {
    return pseg_n.prod(0, start[ind] + 1);
  }

 public:
  ///@param g g[from][j]={to,辺の重み}
  ///@param node_num 頂点の重み
  EulerTour(vector<vector<pair<int, S>>> &g, vector<S> node_num, int root = 0) {
    n = g.size();
    start.resize(n);
    finish.resize(n);
    parent.resize(n);
    euler_tour.resize(n * 2);
    depth.resize(n);
    node_sum.resize(n * 2, es());
    edge_sum.resize(n * 2, et());
    depth_v.resize(n * 2, {INT_MAX / 10, -1});
    rec(root, -1, 0, g);
    for (int i = 0; i < n; i++) {
      euler_tour[start[i]] = i;
      euler_tour[finish[i]] = -i;
      node_sum[start[i]] = node_num[i];
      depth_v[start[i]] = {depth[i], start[i]};
      if (i != root) {
        for (auto nex : g[i]) {
          if (nex.first == parent[i]) {
            edge_sum[start[i]] = nex.second;
            break;
          }
        }
      }
    }
    seg_e = segtree<T, opt, et>(edge_sum);
    seg_n = segtree<S, ops, es>(node_sum);

    for (int i = 0; i < n; i++) {
      node_sum[finish[i]] = invs(node_num[i]);
      if (i != root) {
        for (auto nex : g[i]) {
          if (nex.first == parent[i]) {
            edge_sum[finish[i]] = invt(nex.second);
            depth_v[finish[i]] = {depth[parent[-euler_tour[finish[i]]]], finish[i]};
            break;
          }
        }
      }
    }
    pseg_e = segtree<T, opt, et>(edge_sum);
    pseg_n = segtree<S, ops, es>(node_sum);
    sp_d = SparseTable<SD, f>(depth_v);
  }

  S get_node(int i) {
    return seg_n.get(start[i]);
  }

  T subtree_edge(int ind) {
    return seg_e.prod(start[ind] + 1, finish[ind]);
  }

  S subtree_node(int ind) {
    return seg_n.prod(start[ind], finish[ind]);
  }

  T path_edge(int i, int j) {
    int l = lca(i, j);
    return root_path_edge(i) + root_path_edge(j) - 2 * root_path_edge(l);
  }

  S path_node(int i, int j) {
    int l = lca(i, j);
    return root_path_node(i) + root_path_node(j) - 2 * root_path_node(l);
  }

  void update_weight_edge(int i, int j, T x) {
    if (i != parent[j]) swap(i, j);
    seg_e.set(start[j], x);
    pseg_e.set(start[j], x);
    pseg_e.set(finish[j], invt(x));
  }

  void update_weight_node(int i, T x) {
    seg_n.set(start[i], x);
    pseg_n.set(start[i], x);
    pseg_n.set(finish[i], invt(x));
  }

  int lca(int i, int j) {
    int res;
    if (start[i] > start[j]) swap(i, j);
    SD id = sp_d.query(start[i], start[j] + 1);
    if (euler_tour[id.ind] < 0) {
      res = parent[-euler_tour[id.ind]];
    } else {
      res = euler_tour[id.ind];
    }
    return res;
  }
};
#line 5 "verify/LibraryChecker/tree/VertexAddSubtreeSum.test.cpp"

using S = long long;
S op(S a, S b) {
  return a + b;
}

S e() {
  return 0;
}

S inv(S a) {
  return -a;
}

int main() {
  cin.tie(0)->sync_with_stdio(0);
  int n, q;
  cin >> n >> q;
  vector<vector<pair<int, S>>> g(n);
  vector<S> nodew(n);
  for (int i = 0; i < n; i++) {
    cin >> nodew[i];
  }
  for (int i = 1; i < n; i++) {
    int v;
    cin >> v;
    g[i].push_back({v, 0});
    g[v].push_back({i, 0});
  }

  EulerTour<S, op, e, inv, S, op, e, inv> t(g, nodew);

  for (int i = 0; i < q; i++) {
    int com;
    cin >> com;
    if (com == 0) {
      int u, x;
      cin >> u >> x;
      t.update_weight_node(u, t.get_node(u) + x);
    } else if (com == 1) {
      int u;
      cin >> u;
      cout << t.subtree_node(u) << "\n";
    }
  }
}

Back to top page