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

#include <algorithm>
#include <assert.h>
#include <cmath>
#include <iostream>
#include <stack>
#include <tuple>
#include <vector>

using namespace std;

struct LeastCommonAncestor {
public:
  LeastCommonAncestor(int n, int root = 0)
      : root(root), graph(n), parent(ilogb(n) + 1, vector<int>(n)), depth(n) {}

  void add_edge(int a, int b) {
    graph[a].push_back(b);
    graph[b].push_back(a);
  }

  void init() {
    // dfs(root, -1, 0);

    stack<tuple<int, int, int>> st;
    st.emplace(root, -1, 0);

    while (!st.empty()) {
      int v, p, d;
      tie(v, p, d) = st.top();
      st.pop();

      parent[0][v] = p;
      depth[v] = d;
      for (int to : graph[v]) {
        if (to != p) {
          // dfs(to, v, d + 1);
          st.emplace(to, v, d + 1);
        }
      }
    }

    for (int k = 0; k + 1 < parent.size(); ++k) {
      for (int v = 0; v < graph.size(); ++v) {
        if (parent[k][v] < 0) {
          parent[k + 1][v] = -1;
        } else {
          parent[k + 1][v] = parent[k][parent[k][v]];
        }
      }
    }
  }

  int query(int u, int v) {
    if (depth[u] > depth[v]) {
      swap(u, v);
    }

    for (int k = 0; k < parent.size(); ++k) {
      if ((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }

    if (u == v) {
      return u;
    }

    for (int k = parent.size() - 1; k >= 0; k--) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }

    return parent[0][u];
  }

private:
  void dfs(int v, int p, int d) {
    parent[0][v] = p;
    depth[v] = d;
    for (int to : graph[v]) {
      if (to != p) {
        dfs(to, v, d + 1);
      }
    }
  }

private:
  int root;
  vector<vector<int>> graph;
  vector<vector<int>> parent;
  vector<int> depth;
};

int n, k;

struct Edge {
  Edge() = default;

  Edge(int to, int rev, int64_t cost) : to(to), rev(rev), cost(cost) {}

  int to, rev;
  int64_t cost;
};

struct Query {
  int com, x, y;
};

vector<vector<Edge>> tree;
vector<Edge> sub_tree;

vector<int64_t> lazy;
vector<int64_t> sum;

inline void add_edge(vector<vector<Edge>> &graph, int from, int to,
                     int64_t cost) {
  graph[from].emplace_back(to, graph[to].size(), cost);
  graph[to].emplace_back(from, graph[from].size() - 1, cost);
}

void set_sub_tree(const vector<int> &use) {
  fill(sum.begin(), sum.end(), 0);

  // dfs(use, 0, -1, -1, 0, 0, 0);
  static stack<tuple<int, int, int, int64_t, int64_t, int64_t>> st;
  st.emplace(n / 2, -1, -1, 0, 0, 0);

  while (!st.empty()) {
    int node, par, prev;
    int64_t s, e1_cost, e2_cost;

    tie(node, par, prev, s, e1_cost, e2_cost) = st.top();
    st.pop();

    if (use[node] && prev != -1) {
      if (prev == par) {
        sub_tree[node].to = prev;
        sub_tree[node].cost = e1_cost;
      } else {
        sub_tree[par].to = prev;
        sub_tree[par].cost = e1_cost;

        sub_tree[node].to = par;
        sub_tree[node].cost = e2_cost;
        sum[par] = s;
      }
    }

    if (use[node]) {
      for (auto &e : tree[node]) {
        if (e.to == par) {
          continue;
        }

        int64_t cost = 0;
        st.emplace(e.to, node, node, cost, e.cost, e.cost);
      }
    } else {
      for (auto &e : tree[node]) {
        if (e.to == par) {
          continue;
        }

        int64_t cost = s + (use[e.to] ? 0 : (e.cost % k ? e.cost : 0));
        st.emplace(e.to, node, prev, cost, e1_cost, e.cost);
      }
    }
  }
}

inline void add(int x, int d) { lazy[x] += d; }

inline int64_t send(int s, int t, int u) {
  int64_t res = 0;

  while (s != u) {
    auto &e = sub_tree[s];

    int64_t cost = e.cost + lazy[s] + lazy[e.to];

    if (cost % k == 0) {
      cost = 0;
    }

    cost += sum[s];

    res += cost;

    s = e.to;
  }

  while (t != u) {
    auto &e = sub_tree[t];

    int64_t cost = e.cost + lazy[t] + lazy[e.to];

    if (cost % k == 0) {
      cost = 0;
    }

    cost += sum[t];

    res += cost;

    t = e.to;
  }

  return res;
}

inline void propagate() {
  for (int i = 0; i < n; ++i) {
    if (lazy[i] != 0) {
      for (auto &e : tree[i]) {
        e.cost += lazy[i];
        tree[e.to][e.rev].cost += lazy[i];
      }
    }
  }

  fill(lazy.begin(), lazy.end(), 0);
}

int main() {
  scanf("%d%d", &n, &k);

  tree.assign(n, {});
  sub_tree.assign(n, {});
  lazy.assign(n, 0);
  sum.assign(n, 0);

  LeastCommonAncestor lca(n, n / 2);

  for (int i = 0; i < n - 1; ++i) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);

    add_edge(tree, a, b, c);
    lca.add_edge(a, b);
  }

  lca.init();

  int q;
  scanf("%d", &q);

  assert(n <= 100000 && q <= 100000);

  vector<Query> query(q);

  for (int i = 0; i < q; ++i) {
    char com[5];
    scanf("%s%d%d", com, &query[i].x, &query[i].y);

    query[i].com = (com[0] == 's');
  }

  int bs = 3 * sqrt(q);
  int bn = (q + bs - 1) / bs;

  vector<int> lca_res(q);
  vector<int> use(n);

  for (int i = 0; i < bn; ++i) {
    fill(use.begin(), use.end(), 0);

    for (int j = i * bs; j < (i + 1) * bs && j < q; ++j) {
      if (query[j].com == 1) {
        use[query[j].x] = use[query[j].y] = true;
        lca_res[j] = lca.query(query[j].x, query[j].y);
        use[lca_res[j]] = true;
      } else {
        use[query[j].x] = true;
      }
    }

    set_sub_tree(use);

    for (int j = i * bs; j < (i + 1) * bs && j < q; ++j) {
      if (query[j].com == 0) {
        add(query[j].x, query[j].y);
      } else {
        printf("%lld\\n", send(query[j].x, query[j].y, lca_res[j]));
      }
    }

    propagate();
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <assert.h>
#include <cmath>
#include <iostream>
#include <stack>
#include <tuple>
#include <vector>

using namespace std;

struct LeastCommonAncestor {
public:
  LeastCommonAncestor(int n, int root = 0)
      : root(root), graph(n), parent(ilogb(n) + 1, vector<int>(n)), depth(n) {}

  void add_edge(int a, int b) {
    graph[a].push_back(b);
    graph[b].push_back(a);
  }

  void init() {
    // dfs(root, -1, 0);

    stack<tuple<int, int, int>> st;
    st.emplace(root, -1, 0);

    while (!st.empty()) {
      int v, p, d;
      tie(v, p, d) = st.top();
      st.pop();

      parent[0][v] = p;
      depth[v] = d;
      for (int to : graph[v]) {
        if (to != p) {
          // dfs(to, v, d + 1);
          st.emplace(to, v, d + 1);
        }
      }
    }

    for (int k = 0; k + 1 < parent.size(); ++k) {
      for (int v = 0; v < graph.size(); ++v) {
        if (parent[k][v] < 0) {
          parent[k + 1][v] = -1;
        } else {
          parent[k + 1][v] = parent[k][parent[k][v]];
        }
      }
    }
  }

  int query(int u, int v) {
    if (depth[u] > depth[v]) {
      swap(u, v);
    }

    for (int k = 0; k < parent.size(); ++k) {
      if ((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }

    if (u == v) {
      return u;
    }

    for (int k = parent.size() - 1; k >= 0; k--) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }

    return parent[0][u];
  }

private:
  void dfs(int v, int p, int d) {
    parent[0][v] = p;
    depth[v] = d;
    for (int to : graph[v]) {
      if (to != p) {
        dfs(to, v, d + 1);
      }
    }
  }

private:
  int root;
  vector<vector<int>> graph;
  vector<vector<int>> parent;
  vector<int> depth;
};

int n, k;

struct Edge {
  Edge() = default;

  Edge(int to, int rev, int64_t cost) : to(to), rev(rev), cost(cost) {}

  int to, rev;
  int64_t cost;
};

struct Query {
  int com, x, y;
};

vector<vector<Edge>> tree;
vector<Edge> sub_tree;

vector<int64_t> lazy;
vector<int64_t> sum;

inline void add_edge(vector<vector<Edge>> &graph, int from, int to,
                     int64_t cost) {
  graph[from].emplace_back(to, graph[to].size(), cost);
  graph[to].emplace_back(from, graph[from].size() - 1, cost);
}

void set_sub_tree(const vector<int> &use) {
  fill(sum.begin(), sum.end(), 0);

  // dfs(use, 0, -1, -1, 0, 0, 0);
  static stack<tuple<int, int, int, int64_t, int64_t, int64_t>> st;
  st.emplace(n / 2, -1, -1, 0, 0, 0);

  while (!st.empty()) {
    int node, par, prev;
    int64_t s, e1_cost, e2_cost;

    tie(node, par, prev, s, e1_cost, e2_cost) = st.top();
    st.pop();

    if (use[node] && prev != -1) {
      if (prev == par) {
        sub_tree[node].to = prev;
        sub_tree[node].cost = e1_cost;
      } else {
        sub_tree[par].to = prev;
        sub_tree[par].cost = e1_cost;

        sub_tree[node].to = par;
        sub_tree[node].cost = e2_cost;
        sum[par] = s;
      }
    }

    if (use[node]) {
      for (auto &e : tree[node]) {
        if (e.to == par) {
          continue;
        }

        int64_t cost = 0;
        st.emplace(e.to, node, node, cost, e.cost, e.cost);
      }
    } else {
      for (auto &e : tree[node]) {
        if (e.to == par) {
          continue;
        }

        int64_t cost = s + (use[e.to] ? 0 : (e.cost % k ? e.cost : 0));
        st.emplace(e.to, node, prev, cost, e1_cost, e.cost);
      }
    }
  }
}

inline void add(int x, int d) { lazy[x] += d; }

inline int64_t send(int s, int t, int u) {
  int64_t res = 0;

  while (s != u) {
    auto &e = sub_tree[s];

    int64_t cost = e.cost + lazy[s] + lazy[e.to];

    if (cost % k == 0) {
      cost = 0;
    }

    cost += sum[s];

    res += cost;

    s = e.to;
  }

  while (t != u) {
    auto &e = sub_tree[t];

    int64_t cost = e.cost + lazy[t] + lazy[e.to];

    if (cost % k == 0) {
      cost = 0;
    }

    cost += sum[t];

    res += cost;

    t = e.to;
  }

  return res;
}

inline void propagate() {
  for (int i = 0; i < n; ++i) {
    for (auto &e : tree[i]) {
      e.cost += lazy[i];
      tree[e.to][e.rev].cost += lazy[i];
    }
  }

  fill(lazy.begin(), lazy.end(), 0);
}

int main() {
  scanf("%d%d", &n, &k);

  tree.assign(n, {});
  sub_tree.assign(n, {});
  lazy.assign(n, 0);
  sum.assign(n, 0);

  LeastCommonAncestor lca(n, n / 2);

  for (int i = 0; i < n - 1; ++i) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);

    add_edge(tree, a, b, c);
    lca.add_edge(a, b);
  }

  lca.init();

  int q;
  scanf("%d", &q);

  assert(n <= 100000 && q <= 100000);

  vector<Query> query(q);

  for (int i = 0; i < q; ++i) {
    char com[5];
    scanf("%s%d%d", com, &query[i].x, &query[i].y);

    query[i].com = (com[0] == 's');
  }

  int bs = 3 * sqrt(q);
  int bn = (q + bs - 1) / bs;

  vector<int> lca_res(q);
  vector<int> use(n);

  for (int i = 0; i < bn; ++i) {
    fill(use.begin(), use.end(), 0);

    for (int j = i * bs; j < (i + 1) * bs && j < q; ++j) {
      if (query[j].com == 1) {
        use[query[j].x] = use[query[j].y] = true;
        lca_res[j] = lca.query(query[j].x, query[j].y);
        use[lca_res[j]] = true;
      } else {
        use[query[j].x] = true;
      }
    }

    set_sub_tree(use);

    for (int j = i * bs; j < (i + 1) * bs && j < q; ++j) {
      if (query[j].com == 0) {
        add(query[j].x, query[j].y);
      } else {
        printf("%lld\\n", send(query[j].x, query[j].y, lca_res[j]));
      }
    }

    propagate();
  }
}