Problem p02375 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

template <typename T> constexpr function<T(T, T)> ex_assign(T initE) {
  return [&initE](const T &x, const T &y) -> const T {
    return (y == initE) ? x : y;
  };
}
template <typename T> constexpr function<T(T, int)> ex_multiplies(T initE) {
  return [&initE](const T &x, size_t y) -> const T {
    return (x == initE) ? x : x * y;
  };
}
template <typename T, typename E> struct LazySegmentTree {
  using F = function<T(T, T)>;
  using G = function<T(T, E)>;
  using H = function<E(E, E)>;
  using P = function<E(E, int)>;

  int n;
  T initT;
  E initE;
  F f;
  G g;
  H h;
  P p;
  vector<T> data;
  vector<E> lazy;

  LazySegmentTree() {}
  LazySegmentTree(
      int n_, F f, G g, H h, T initT = INT_MAX, E initE = INT_MAX,
      vector<T> v = vector<T>(),
      P p =
          [](E a, size_t b) {
            ++b;
            return a;
          })
      : f(f), g(g), h(h), initT(initT), initE(initE), p(p) {
    n = 1;
    while (n < n_)
      n *= 2;
    data.assign(n * 2 - 1, initT);
    lazy.assign(n * 2 - 1, initE);
    if (n_ == v.size())
      build(v);
  }
  void build(vector<T> v) {
    for (int i = 0; i < v.size(); ++i)
      data[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; --i) {
      data[i] = f(data[i * 2 + 1], data[i * 2 + 2]);
    }
  }
  inline void eval(int len, int k) {
    if (lazy[k] == initE)
      return;
    if (k * 2 + 1 < n * 2 - 1) {
      lazy[k * 2 + 1] = h(lazy[k * 2 + 1], lazy[k]);
      lazy[k * 2 + 2] = h(lazy[k * 2 + 2], lazy[k]);
    }
    data[k] = g(data[k], p(lazy[k], len));
    lazy[k] = initE;
  }
  T update(int a, int b, E x, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    eval(r - l, k);
    if (b <= l || r <= a)
      return data[k];
    if (a <= l && r <= b) {
      lazy[k] = h(lazy[k], x);
      return g(data[k], p(lazy[k], r - l));
    }
    T vl = update(a, b, x, k * 2 + 1, l, (l + r) / 2);
    T vr = update(a, b, x, k * 2 + 2, (l + r) / 2, r);
    return data[k] = f(vl, vr);
  }
  T query(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    eval(r - l, k);
    if (b <= l || r <= a)
      return initT;
    if (a <= l && r <= b)
      return data[k];
    T vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
    T vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(vl, vr);
  }
};
//
// e.g.
//      |     f     |   g   |   h   |   initT   | initE |   p
// -----+-----------+-------+-------+-----------+-------+--------
//  RMQ | min / max |       |       | INF / 0   |       | (mult)
//  RSQ | plus      |       |       | 0         |       | (mult)
//  RUQ |           | (ass) | (ass) |           | INF   | (mult)
//  RAQ |           | plus  | plus  |           | 0     | (mult)
//
//  (ass)  := ex_assign;
//  (mult) := ex_multiplies;
//

struct HeavyLightDecomposition {
  using P = pair<int, int>;
  using F = function<void(int, int)>;

  int n, time;
  vector<vector<int>> g;
  vector<int> num, next, subSize, head, par, depth, inv, chain;

  HeavyLightDecomposition() {}
  HeavyLightDecomposition(int n)
      : n(n), time(0), g(n), num(n, -1), next(n, -1), subSize(n, 1), head(n),
        par(n), depth(n), inv(n), chain(n) {}

  void build(vector<int> roots = vector<int>(1, 0)) {
    int c = 0;
    for (auto &&r : roots) {
      dfs(r);
      bfs(r, c++);
    }
  }
  void addEdge(int s, int t) {
    g[s].push_back(t);
    g[t].push_back(s);
  }
  void dfs(int r) {
    stack<P, deque<P>> s;
    par[r] = -1;
    depth[r] = 0;
    s.push({r, 0});
    while (!s.empty()) {
      int u = s.top().first;
      int &i = s.top().second;
      int sz = g[u].size();
      if (i < sz) {
        int v = g[u][i++];
        if (v == par[u])
          continue;
        par[v] = u;
        depth[v] = depth[u] + 1;
        s.push({v, 0});
      } else {
        s.pop();
        int mx = 0;
        for (auto &&v : g[u]) {
          if (v == par[u])
            continue;
          subSize[u] += subSize[v];
          if (mx < subSize[v])
            mx = subSize[v], next[u] = v;
        }
      }
    }
  }
  void bfs(int r, int c) {
    int &k = time;
    queue<int> q({r});
    while (!q.empty()) {
      int h = q.front();
      q.pop();
      for (int u = h; u != -1; u = next[u]) {
        chain[u] = c;
        num[u] = k++;
        inv[num[u]] = u;
        head[u] = h;
        for (auto &&v : g[u]) {
          if (v != par[u] && v != next[u])
            q.push(v);
        }
      }
    }
  }
  int for_each_vertex(int u, int v, const F &f) {
    while (true) {
      if (num[u] > num[v])
        swap(u, v);
      f(max(num[head[v]], num[u]), num[v]);
      if (head[u] != head[v]) {
        v = par[head[v]];
      } else {
        break;
      }
    }
  }
  int for_each_edge(int u, int v, const F &f) {
    while (true) {
      if (num[u] > num[v])
        swap(u, v);
      if (head[u] != head[v]) {
        f(num[head[v]], num[v]);
        v = par[head[v]];
      } else {
        if (u != v)
          f(num[u] + 1, num[v]);
        break;
      }
    }
  }
  int lowestCommonAncestor(int u, int v) {
    while (true) {
      if (num[u] > num[v])
        swap(u, v);
      if (head[u] == head[v])
        return u;
      v = par[head[v]];
    }
  }
  int distance(int u, int v) {
    int lca = lowestCommonAncestor(u, v);
    return depth[u] + depth[v] - 2 * depth[lca];
  }
};

int main() {
  int n;
  cin >> n;
  HeavyLightDecomposition hld(n);
  for (int i = 0; i < n; ++i) {
    int k;
    cin >> k;
    while (k--) {
      int c;
      cin >> c;
      hld.addEdge(i, c);
    }
  }
  hld.build();

  const ll initT = 0;
  const ll initE = 0;
  const LazySegmentTree<ll, ll>::F f = plus<ll>();
  const LazySegmentTree<ll, ll>::G g = plus<ll>();
  const LazySegmentTree<ll, ll>::H h = plus<ll>();
  const LazySegmentTree<ll, ll>::P p = ex_multiplies<ll>(initE);

  LazySegmentTree<ll, ll> seg(n, f, g, h, initT, initE, vector<ll>(), p);
  auto add = [&](int v, int w) {
    hld.for_each_edge(0, v, [&](int l, int r) { seg.update(l, r + 1, w); });
  };
  auto getSum = [&](int u) {
    ll res = 0;
    hld.for_each_edge(0, u, [&](int l, int r) { res += seg.query(l, r + 1); });
    return res;
  };

  int q;
  cin >> q;
  while (q--) {
    int com;
    cin >> com;
    if (com) {
      int u;
      cin >> u;
      cout << getSum(u) << endl;
    } else {
      int v, w;
      cin >> v >> w;
      add(v, w);
    }
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

template <typename T> struct KdSegmentTree {
  using P = pair<int, int>;
  int n;
  vector<vector<T>> data;

  KdSegmentTree() {}
  KdSegmentTree(int n_, vector<T> v) {
    n = 1;
    while (n < n_)
      n *= 2;
    data.resize(2 * n - 1);
    build(v);
  }
  void build(vector<T> v) {
    for (int i = 0; i < v.size(); ++i)
      data[i + n - 1].push_back(v[i]);
    for (int i = n - 2; i >= 0; --i) {
      int lch = i * 2 + 1;
      int rch = i * 2 + 2;
      data[i].resize(data[lch].size() + data[rch].size());
      merge(begin(data[lch]), end(data[lch]), begin(data[rch]), end(data[rch]),
            begin(data[i]));
    }
  }
  void update(int a, int b, int sy, int ty, int x, vector<int> &v, int k = 0,
              int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    if (b <= l || r <= a)
      return;
    if (a <= l && r <= b) {
      auto lb = lower_bound(begin(data[k]), end(data[k]), P{sy, -1});
      auto ub = upper_bound(begin(data[k]), end(data[k]), P{ty, n});
      for (auto itr = lb; itr != ub; ++itr) {
        if (itr->second == 0)
          continue;
        v[itr->second] += x;
      }
      return;
    }
    update(a, b, sy, ty, x, v, k * 2 + 1, l, (l + r) / 2);
    update(a, b, sy, ty, x, v, k * 2 + 2, (l + r) / 2, r);
    return;
  }
  int query(int a, int b, int sy, int ty, vector<int> &v, int k = 0, int l = 0,
            int r = -1) {
    if (r < 0)
      r = n;
    if (b <= l || r <= a)
      return 0;
    if (a <= l && r <= b) {
      auto lb = lower_bound(begin(data[k]), end(data[k]), P{sy, -1});
      auto ub = upper_bound(begin(data[k]), end(data[k]), P{ty, n});
      int res = 0;
      for (auto itr = lb; itr != ub; ++itr) {
        res += v[itr->second];
      }
      return res;
    }
    int vl = query(a, b, sy, ty, v, k * 2 + 1, l, (l + r) / 2);
    int vr = query(a, b, sy, ty, v, k * 2 + 2, (l + r) / 2, r);
    return vl + vr;
  }
};

struct KdTree {
  using P = pair<int, int>;

  int n;
  vector<int> xs;
  vector<P> ys;
  vector<pair<int, P>> ps;
  KdSegmentTree<P> seg;
  vector<int> es;

  KdTree(int n) : n(n), xs(n), ys(n), ps(n), es(n, 0) {}

  void init() {
    sort(begin(ps), end(ps));
    for (int i = 0; i < n; ++i)
      tie(xs[i], ys[i]) = ps[i];
    seg = KdSegmentTree<P>(n, ys);
  }
  void addPoint(int x, int y, int i) { ps[i] = {x, {y, i}}; }
  void update(int sx, int tx, int sy, int ty, int w) {
    int l = distance(begin(xs), lower_bound(begin(xs), end(xs), sx));
    int r = distance(begin(xs), upper_bound(begin(xs), end(xs), tx));
    seg.update(l, r, sy, ty, w, es);
  }
  int query(int sx, int tx, int sy, int ty) {
    int l = distance(begin(xs), lower_bound(begin(xs), end(xs), sx));
    int r = distance(begin(xs), upper_bound(begin(xs), end(xs), tx));
    return seg.query(l, r, sy, ty, es);
  }
};

struct Tree {
  using P = pair<int, int>;

  int n;
  int time, cnt;
  vector<vector<int>> g;
  vector<int> numX, numY;
  vector<P> depth;
  vector<int> eulerTour;
  // SegmentTree<P> seg;

  Tree(int n)
      : n(n), time(0), cnt(0), g(n), numX(n), numY(n), depth(n * 2 - 1),
        eulerTour(n * 2 - 1) {}

  void init(int r) {
    dfs(r, r, 0);
    // SegmentTree<P>::F f = [](P a, P b) { return min(a, b); };
    // seg = SegmentTree<P>(n * 2 - 1, f, {INT_MAX, -1}, depth);
  }
  void addEdge(int s, int t) {
    g[s].push_back(t);
    g[t].push_back(s);
  }
  void dfs(int u, int p, int d) {
    numX[u] = time;
    eulerTour[time] = u;
    depth[time] = {d, time};
    ++time;
    for (auto &&v : g[u]) {
      if (v != p) {
        dfs(v, u, d + 1);
        eulerTour[time] = u;
        depth[time] = {d, time};
        ++time;
      }
    }
    numY[u] = time;
  }
  // int lowestCommonAncestor(int u, int v) {
  //   int i = num[u];
  //   int j = num[v];
  //   if (i > j) swap(i, j);
  //   return eulerTour[seg.query(i, j + 1).second];
  // }
};

int main() {
  int n;
  scanf("%d", &n);
  Tree t(n);
  for (int i = 0; i < n; ++i) {
    int k;
    scanf("%d", &k);
    while (k--) {
      int c;
      scanf("%d", &c);
      t.addEdge(i, c);
    }
  }
  t.init(0);

  map<int, pair<int, int>> mp;
  KdTree kd(n);
  for (int i = 0; i < n; ++i) {
    int x = t.numX[i];
    int y = t.numY[i];
    mp[i] = {x, y};
    kd.addPoint(x, y, i);
  }
  kd.init();
  int q;
  scanf("%d", &q);
  while (q--) {
    int com;
    scanf("%d", &com);
    if (com) {
      int u;
      scanf("%d", &u);
      int sx, sy;
      tie(sx, sy) = mp[0];
      int tx, ty;
      tie(tx, ty) = mp[u];
      swap(sy, ty);
      printf("%d\\n", kd.query(sx, tx, sy, ty));
    } else {
      int v, w;
      scanf("%d%d", &v, &w);
      int sx, sy;
      tie(sx, sy) = mp[0];
      int tx, ty;
      tie(tx, ty) = mp[v];
      swap(sy, ty);
      kd.update(sx, tx, sy, ty, w);
    }
  }
  return 0;
}