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

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

using Vertex = int;
using Weight = long long;
struct Edge {
  Vertex from, to;
  Weight weight;
};
using Graph = vector<vector<Edge>>;

using Size = int;

template <typename T> class SegmentTree {
public:
  SegmentTree(Size N, T K) : N(pow(N)), node(2 * pow(N) - 1), K(K) {}
  void add(Size i, T v) {
    i += N - 1;
    node[i] += v;
    node[i = parent(i)] = 0;
    if (node[lchild(i)] % K)
      node[i] += node[lchild(i)];
    if (node[rchild(i)] % K)
      node[i] += node[rchild(i)];
    while (i) {
      i = parent(i);
      node[i] = node[lchild(i)] + node[rchild(i)];
    }
  }
  T sum(Size a, Size b) const { return sum(a, b, 0, 0, N); }
  T get(Size i) const { return node[i + N - 1]; }

private:
  constexpr Size pow(Size n) const {
    Size r = 1;
    while (r < n)
      r <<= 1;
    return r;
  }
  inline Size parent(Size n) const { return (n - 1) / 2; }
  inline Size lchild(Size n) const { return 2 * n + 1; }
  inline Size rchild(Size n) const { return 2 * n + 2; }
  inline bool leaf(Size n) const { return N - 1 <= n; }
  T sum(Size a, Size b, Size k, Size l, Size r) const {
    if (r <= a || b <= l)
      return 0;
    if (a <= l && r <= b) {
      if (k < N - 1)
        return node[k];
      return node[k] % K ? node[k] : 0;
    }
    auto m = (l + r) / 2;
    return sum(a, b, lchild(k), l, m) + sum(a, b, rchild(k), m, r);
  }
  const Size N;
  const T K;
  vector<T> node;
};

class CentroidPathDecomposition {
public:
  CentroidPathDecomposition(const Graph &G, Vertex root, Weight K)
      : head(G.size(), NIL), parent(G.size(), NIL), depth(G.size(), 0),
        subsize(G.size(), 1), G(G), K(K), segtree(G.size(), K),
        position(G.size()), child(G.size(), NIL), cost(G.size(), 0) {
    function<Size(Vertex)> dfs = [&](Vertex cur) {
      for (const auto &e : G[cur])
        if (e.to != parent[cur]) {
          depth[e.to] = depth[cur] + 1;
          parent[e.to] = cur;
          subsize[cur] += dfs(e.to);
        }
      return subsize[cur];
    };

    function<Size(Vertex, Size)> build = [&](Vertex cur, Size i) {
      Vertex nex = NIL;
      Size size = 0;
      Weight weight = 0;
      for (const auto &e : G[cur])
        if (e.to != parent[cur])
          if (chmax(size, subsize[e.to]) == subsize[e.to])
            nex = e.to, weight = e.weight;
      if (nex == NIL)
        return i;
      head[nex] = head[cur];
      child[cur] = nex;
      segtree.add(position[nex] = i++, weight);
      i = build(nex, i);
      for (const auto &e : G[cur])
        if (e.to != parent[cur])
          if (e.to != nex) {
            head[e.to] = e.to;
            segtree.add(position[e.to] = i++, e.weight);
            i = build(e.to, i);
          }
      return i;
    };

    dfs(root);
    build(head[root] = root, 1);
  }

  void add(Vertex x, Weight d) {
    cost[x] += d;
    for (auto v : {child[x], x}) {
      if (v == NIL)
        continue;
      if (parent[v] == NIL)
        continue;
      if (head[v] != head[parent[v]])
        continue;
      segtree.add(position[v], d);
    }
  }

  Weight send(Vertex u, Vertex v) {
    Weight res = 0;
    while (head[u] != head[v]) {
      if (depth[head[u]] > depth[head[v]])
        swap(u, v);
      res += segtree.sum(position[head[v]] + 1, position[v] + 1);
      auto w = segtree.get(position[head[v]]);
      w += cost[head[v]];
      w += cost[parent[head[v]]];
      if (w % K)
        res += w;
      v = parent[head[v]];
    }
    if (depth[u] > depth[v])
      swap(u, v);
    if (u != v)
      res += segtree.sum(position[u] + 1, position[v] + 1);
    return res;
  }

private:
  template <typename T> inline T chmax(T &l, const T &r) {
    return l = std::max(l, r);
  };

  const Size NIL = -1;
  const Graph G;
  const Weight K;

  vector<Vertex> head, parent;
  vector<Size> depth, subsize;
  SegmentTree<Weight> segtree;
  vector<Size> position;
  vector<Vertex> child;
  vector<Weight> cost;
};

int main() {
  int N, K;
  cin >> N >> K;

  Graph G(N);
  for (auto i = 0; i < N - 1; ++i) {
    Vertex a, b;
    Weight c;
    cin >> a >> b >> c;
    G[a].push_back({a, b, c});
    G[b].push_back({b, a, c});
  }

  CentroidPathDecomposition T(G, 0, K);

  int Q;
  cin >> Q;
  while (Q--) {
    string s;
    cin >> s;
    if (s == "add") {
      Vertex x;
      Weight d;
      cin >> x >> d;
      T.add(x, d);
    } else {
      Vertex s, t;
      cin >> s >> t;
      cout << T.send(s, t) << endl;
    }
  }
}

Error: TLE

Fixed submission:

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

using Vertex = int;
using Weight = long long;
struct Edge {
  Vertex from, to;
  Weight weight;
};
using Graph = vector<vector<Edge>>;

using Size = int;

template <typename T> class SegmentTree {
public:
  SegmentTree(Size N, T K) : N(pow(N)), node(2 * pow(N) - 1), K(K) {}
  void add(Size i, T v) {
    i += N - 1;
    node[i] += v;
    node[i = parent(i)] = 0;
    if (node[lchild(i)] % K)
      node[i] += node[lchild(i)];
    if (node[rchild(i)] % K)
      node[i] += node[rchild(i)];
    while (i) {
      i = parent(i);
      node[i] = node[lchild(i)] + node[rchild(i)];
    }
  }
  T sum(Size a, Size b) const { return sum(a, b, 0, 0, N); }

private:
  constexpr Size pow(Size n) const {
    Size r = 1;
    while (r < n)
      r <<= 1;
    return r;
  }
  inline Size parent(Size n) const { return (n - 1) / 2; }
  inline Size lchild(Size n) const { return 2 * n + 1; }
  inline Size rchild(Size n) const { return 2 * n + 2; }
  inline bool leaf(Size n) const { return N - 1 <= n; }
  T sum(Size a, Size b, Size k, Size l, Size r) const {
    if (r <= a || b <= l)
      return 0;
    if (a <= l && r <= b) {
      if (k < N - 1)
        return node[k];
      return node[k] % K ? node[k] : 0;
    }
    auto m = (l + r) / 2;
    return sum(a, b, lchild(k), l, m) + sum(a, b, rchild(k), m, r);
  }
  const Size N;
  const T K;
  vector<T> node;
};

class CentroidPathDecomposition {
public:
  CentroidPathDecomposition(const Graph &G, Vertex root, Weight K)
      : head(G.size(), NIL), parent(G.size(), NIL), depth(G.size(), 0),
        subsize(G.size(), 1), G(G), segtree(G.size(), K), position(G.size()) {
    function<Size(Vertex)> dfs = [&](Vertex cur) {
      for (const auto &e : G[cur])
        if (e.to != parent[cur]) {
          depth[e.to] = depth[cur] + 1;
          parent[e.to] = cur;
          subsize[cur] += dfs(e.to);
        }
      return subsize[cur];
    };

    function<Size(Vertex, Size)> build = [&](Vertex cur, Size i) {
      position[cur] = i++;
      Vertex nex = NIL;
      Size size = 0;
      Weight weight = 0;
      for (const auto &e : G[cur])
        if (e.to != parent[cur])
          if (chmax(size, subsize[e.to]) == subsize[e.to])
            nex = e.to, weight = e.weight;
      if (nex == NIL)
        return i;
      head[nex] = head[cur];
      segtree.add(i, weight);
      i = build(nex, i);
      for (const auto &e : G[cur])
        if (e.to != parent[cur])
          if (e.to != nex) {
            head[e.to] = e.to;
            segtree.add(i, e.weight);
            i = build(e.to, i);
          }
      return i;
    };

    dfs(root);
    build(head[root] = root, 0);
  }

  void add(Vertex x, Weight d) {
    for (const auto &e : G[x]) {
      auto u = e.from, v = e.to;
      if (depth[u] > depth[v])
        swap(u, v);
      segtree.add(position[v], d);
    }
  }

  Weight send(Vertex u, Vertex v) {
    Weight res = 0;
    while (head[u] != head[v]) {
      if (depth[head[u]] > depth[head[v]])
        swap(u, v);
      res += segtree.sum(position[head[v]], position[v] + 1);
      v = parent[head[v]];
    }
    if (depth[u] > depth[v])
      swap(u, v);
    if (u != v)
      res += segtree.sum(position[u] + 1, position[v] + 1);
    return res;
  }

private:
  template <typename T> inline T chmax(T &l, const T &r) {
    return l = std::max(l, r);
  };

  const Size NIL = -1;
  const Graph G;

  vector<Vertex> head, parent;
  vector<Size> depth, subsize;
  SegmentTree<Weight> segtree;
  vector<Size> position;
};

int main() {
  int N, K;
  cin >> N >> K;

  Graph G(N);
  for (auto i = 0; i < N - 1; ++i) {
    Vertex a, b;
    Weight c;
    cin >> a >> b >> c;
    G[a].push_back({a, b, c});
    G[b].push_back({b, a, c});
  }

  CentroidPathDecomposition T(G, 0, K);

  int Q;
  cin >> Q;
  while (Q--) {
    string s;
    cin >> s;
    if (s == "add") {
      Vertex x;
      Weight d;
      cin >> x >> d;
      T.add(x, d);
    } else {
      Vertex s, t;
      cin >> s >> t;
      cout << T.send(s, t) << "\\n";
    }
  }
}