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

#include <bits/stdc++.h>

using namespace std;

using int64 = long long;
const int INF = 1 << 30;

struct edge {
  int to;
  int cost;

  pair<int, int> to_adjust, curr_adjust;
};

int N, K, Q;
char S[500][10];
int A[500], B[500];
vector<vector<edge>> g;
int64 weight[100000];
vector<vector<edge>> gg;

struct TreeContraction {
  vector<int> marker, order, parent;

  TreeContraction(int n) {
    g.resize(n);
    destroy();
  }

  void destroy() { marker.assign(g.size(), 0); }

  void add_edge(int idx, const edge &e) { g[idx].emplace_back(e); }

  void build() {
    function<void(int, int)> dfs = [&](int idx, int par) {
      for (auto &e : g[idx]) {
        if (e.to == par)
          continue;
        dfs(e.to, idx);
      }
      parent[idx] = par;
      order.emplace_back(idx);
    };
    parent.resize(g.size());
    dfs(0, -1);
    order.pop_back();
  }

  inline void mark(int x) { marker[x] = 2; }

  inline edge contraction(int idx, int par) {
    if (marker[idx] == 0)
      return (edge){-1};
    for (auto &e : g[idx]) {
      if (e.to == par)
        continue;
      auto ret = contraction(e.to, idx);
      if (ret.to == -1)
        continue;
      ret.cost += e.cost;
      if (ret.to_adjust.first == -1) {
        ret.to_adjust = {idx, e.cost};
      } else if (marker[idx] <= 1) { // to-to_adjust * - -!idx
        auto sum = e.cost + weight[e.to] + weight[idx];
        if (sum % K == 0)
          ret.cost -= sum; // えーわからない
      }
      if (marker[idx] == 1) {
        ret.cost += weight[idx] * 2;
        return (ret);
      } else {
        ret.curr_adjust = {e.to, e.cost};
        gg[idx].emplace_back(ret);
        gg[ret.to].emplace_back(
            (edge){idx, ret.cost, ret.curr_adjust, ret.to_adjust});
      }
    }
    return ((edge){idx, 0, make_pair(-1, -1)});
  };

  inline void contract() {
    for (int i : order)
      if (marker[i] >= 1)
        ++marker[parent[i]];
    gg.clear();
    gg.resize(N);
    contraction(0, -1);
  }
};

int64 cost;

inline bool dfs2(int s, int t, int par) {
  if (s == t) {
    cost += weight[t];
    cost += weight[t];
    return (true);
  }
  for (auto &e : gg[s]) {
    if (e.to == par)
      continue;
    if (dfs2(e.to, t, s)) {
      auto sum = weight[s] + weight[e.curr_adjust.first] + e.curr_adjust.second;
      if (sum % K == 0)
        cost -= sum;
      if (e.to_adjust.first != s) {
        sum = weight[e.to] + weight[e.to_adjust.first] + e.to_adjust.second;
        if (sum % K == 0)
          cost -= sum;
      }
      cost += e.cost;
      cost += weight[s];
      cost += weight[s];
      return (true);
    }
  }
  return (false);
}

int main() {
  scanf("%d %d", &N, &K);
  TreeContraction tree(N);
  for (int i = 1; i < N; i++) {
    int a, b, c;
    scanf("%d %d %d", &a, &b, &c);
    tree.add_edge(a, (edge){b, c});
    tree.add_edge(b, (edge){a, c});
  }
  tree.build();
  scanf("%d", &Q);

  for (int i = 0; i < Q; i += 360) {
    int l = i, r = min(i + 360, Q);
    tree.destroy();
    for (int j = 0; j < r - l; j++) {
      scanf(" %s %d %d", S[j], &A[j], &B[j]);
      if (S[j][0] == 'a')
        tree.mark(A[j]);
      else
        tree.mark(A[j]), tree.mark(B[j]);
    }
    tree.contract();
    for (int j = 0; j < r - l; j++) {
      if (S[j][0] == 'a') {
        weight[A[j]] += B[j];
      } else {
        cost = 0;
        dfs2(B[j], A[j], -1);
        printf("%lld\\n", cost - weight[A[j]] - weight[B[j]]);
      }
    }
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

using int64 = long long;
const int INF = 1 << 30;

struct edge {
  int to;
  int64 cost;

  pair<int, int> to_adjust, curr_adjust;
};

int N, K, Q;
char S[500][10];
int A[500], B[500];
vector<vector<edge>> g;
int64 weight[100000];
vector<vector<edge>> gg;

struct TreeContraction {
  vector<int> marker, order, parent;

  TreeContraction(int n) {
    g.resize(n);
    destroy();
  }

  void destroy() { marker.assign(g.size(), 0); }

  void add_edge(int idx, const edge &e) { g[idx].emplace_back(e); }

  void build() {
    function<void(int, int)> dfs = [&](int idx, int par) {
      for (auto &e : g[idx]) {
        if (e.to == par)
          continue;
        dfs(e.to, idx);
      }
      parent[idx] = par;
      order.emplace_back(idx);
    };
    parent.resize(g.size());
    dfs(0, -1);
    order.pop_back();
  }

  inline void mark(int x) { marker[x] = 2; }

  inline edge contraction(int idx, int par) {
    if (marker[idx] == 0)
      return (edge){-1};
    for (auto &e : g[idx]) {
      if (e.to == par)
        continue;
      auto ret = contraction(e.to, idx);
      if (ret.to == -1)
        continue;
      ret.cost += e.cost;
      if (ret.to_adjust.first == -1) {
        ret.to_adjust = {idx, e.cost};
      } else if (marker[idx] <= 1) {
        auto sum = e.cost + weight[e.to] + weight[idx];
        if (sum % K == 0)
          ret.cost -= sum;
      }
      if (marker[idx] == 1) {
        ret.cost += weight[idx] * 2;
        return (ret);
      } else {
        ret.curr_adjust = {e.to, e.cost};
        gg[idx].emplace_back(ret);
        gg[ret.to].emplace_back(
            (edge){idx, ret.cost, ret.curr_adjust, ret.to_adjust});
      }
    }
    return ((edge){idx, 0, make_pair(-1, -1)});
  };

  inline void contract() {
    for (int i : order)
      if (marker[i] >= 1)
        ++marker[parent[i]];
    gg.clear();
    gg.resize(N);
    contraction(0, -1);
  }
};

int64 cost;
int t;

inline bool dfs2(int s, int par) {
  if (s == t) {
    cost += weight[t];
    cost += weight[t];
    return (true);
  }
  for (auto &e : gg[s]) {
    if (e.to == par)
      continue;
    if (dfs2(e.to, s)) {
      auto sum = weight[s] + weight[e.curr_adjust.first] + e.curr_adjust.second;
      if (sum % K == 0)
        cost -= sum;
      if (e.to_adjust.first != s) {
        sum = weight[e.to] + weight[e.to_adjust.first] + e.to_adjust.second;
        if (sum % K == 0)
          cost -= sum;
      }
      cost += e.cost;
      cost += weight[s];
      cost += weight[s];
      return (true);
    }
  }
  return (false);
}

int main() {
  scanf("%d %d", &N, &K);
  TreeContraction tree(N);
  for (int i = 1; i < N; i++) {
    int a, b, c;
    scanf("%d %d %d", &a, &b, &c);
    tree.add_edge(a, (edge){b, c});
    tree.add_edge(b, (edge){a, c});
  }
  tree.build();
  scanf("%d", &Q);

  for (int i = 0; i < Q; i += 360) {
    int l = i, r = min(i + 360, Q);
    tree.destroy();
    for (int j = 0; j < r - l; j++) {
      scanf(" %s %d %d", S[j], &A[j], &B[j]);
      if (S[j][0] == 'a')
        tree.mark(A[j]);
      else
        tree.mark(A[j]), tree.mark(B[j]);
    }
    tree.contract();
    for (int j = 0; j < r - l; j++) {
      if (S[j][0] == 'a') {
        weight[A[j]] += B[j];
      } else {
        cost = 0;
        t = A[j];
        dfs2(B[j], -1);
        printf("%lld\\n", cost - weight[A[j]] - weight[B[j]]);
      }
    }
  }
}