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

// GEKI TSU RA I

// see https://gist.github.com/LumaKernel/ff55d49ee1af69b7388f15b707e75c15
#include <iostream>
#include <vector>

using namespace std;

typedef long long ll;
typedef pair<int, int> P;
typedef vector<int> VI;
typedef vector<P> VP;

#define omajinai                                                               \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0)

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define RFOR(i, b, a) for (int i = (b)-1; i >= (a); --i)
#define RREP(i, n) RFOR(i, n, 0)

#define LFOR(i, a, b) for (ll i = (a); i < (b); ++i)
#define RLFOR(i, b, a) for (ll i = (b)-1; i >= (a); --i)

#define ALL(a) (a).begin(), (a).end()
#define UNIQUE(a) (a).erase(unique((a).begin(), (a).end()), (a).end())
#define MP make_pair
#define PB push_back
#define EACH(i, c) REP(i, (int)(c).size())
#define EXIST(s, e) ((s).find(e) != (s).end())
#define SORT(c) sort((c).begin(), (c).end())

#define dump(x) cerr << "[L " << __LINE__ << "] " << #x << " = " << (x) << "\\n";
#define dump2(x, y)                                                            \\
  cerr << "[L " << __LINE__ << "] " << #x << " = " << (x) << " , " << #y       \\
       << " = " << (y) << "\\n";

const int INF = 1e9;
const double EPS = 1e-10;
const int _N = 1e5;

int n, k, q;
VP edges[_N];

struct Node {
  ll top = 0, center = 0, bottom = 0;
  Node() {} //
  Node(ll top, ll center, ll bottom)
      : top(top), center(center), bottom(bottom) {}
};

const Node ZERO(-1, 0, 0);

Node merge(Node left, Node right) {
  if (left.top == -1)
    return right;
  if (right.top == -1)
    return left;

  ll center = left.center + right.center;
  ll len = left.bottom + right.top;
  if (len % k != 0)
    center += len; //

  return Node(left.top, center, right.bottom);
}

struct SegmentTree {
  Node ds[_N * 4 + 1];
  int m;
  Node temp;

  void init(int n) {
    m = 1;
    while (m < n)
      m *= 2;
    REP(i, m * 2 - 1)
    ds[i] = ZERO;
  }

  void update(int i, Node x) {
    i += m - 1;
    ds[i] = x;
    while (i >= 1) {
      i = (i - 1) / 2;
      ds[i] = merge(ds[i * 2 + 1], ds[i * 2 + 2]);
    }
  }

  Node get(int i) { return ds[i + m - 1]; }
  Node query(int a, int b) {
    temp = ZERO;
    _query(a, b, 0, 0, m);
    return temp;
  }

  void _query(int a, int b, int k, int l, int r) {
    if (r <= a || b <= l)
      return; //

    if (a <= l && r <= b) {
      temp = merge(temp, ds[k]);
      return;
    }
    _query(a, b, k * 2 + 1, l, (r + l) / 2);
    _query(a, b, k * 2 + 2, (r + l) / 2, r);
  }
};

struct X {};
X x;

int par[_N], head[_N], size[_N], depth[_N];
int indx[_N], id;

int dfs1(int t, int p) {
  par[t] = p;
  if (p != -1)
    depth[t] = depth[p] + 1; //
  int sz = 1;
  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    sz += dfs1(j, t);
  }
  return size[t] = sz;
}

SegmentTree seg;

void dfs2(int t, int p, int d) {
  //
  indx[t] = id;
  seg.update(id, Node(d, 0, 0));
  id++;

  int maxP = -1;
  int maxSz = 0;
  int maxCost = 0;
  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    if (maxSz < size[j]) { //
      maxP = j;
      maxSz = size[j];
      maxCost = edges[t][i].second;
    }
  }
  if (maxP == -1)
    return;
  head[maxP] = head[t];

  dfs2(maxP, t, maxCost); // for SegmentTree

  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    if (j == maxP)
      continue;
    dfs2(j, t, edges[t][i].second);
  }
}

// HL-Decomposition
void decomposite() {
  dfs1(0, -1);
  REP(i, n) { head[i] = i; }
  dfs2(0, -1, 0);
}

int lca(int a, int b) {
  int p = head[a];
  int q = head[b];
  if (p == q)
    return depth[a] < depth[b] ? a : b;
  if (depth[p] < depth[q]) {
    return lca(a, par[q]);
  } else {
    return lca(par[p], b);
  }
}

ll dist(int p, int top) {
  Node d = ZERO;
  while (true) {
    if (head[p] == head[top]) {
      d = merge(seg.query(indx[top], indx[p] + 1), d);
      break;
    }
    d = merge(seg.query(indx[head[p]], indx[p] + 1), d);

    p = par[head[p]];
  }

  return d.center;
}

int main() {
  omajinai;
  cin >> n >> k;
  REP(i, n - 1) {
    int a, b, c;
    cin >> a >> b >> c;
    edges[a].PB(MP(b, c));
    edges[b].PB(MP(a, c));
  }

  seg.init(n);
  decomposite();

  cin >> q;

  REP(i, q) {
    string s;
    cin >> s;
    if (s[0] == 'a') { // add
      int x, d;
      cin >> x >> d;

      Node one = seg.get(indx[x]);
      one.top += d;
      one.bottom += d;
      seg.update(indx[x], one);
    } else { // send
      int s, t;
      cin >> s >> t;
      int l = lca(s, t);
      ll d = dist(s, l) + dist(t, l);
      cout << d << endl;
    }
  }
}

Error: TLE

Fixed submission:

// GEKI TSU RA I

// see https://gist.github.com/LumaKernel/ff55d49ee1af69b7388f15b707e75c15
#include <iostream>
#include <vector>

using namespace std;

typedef long long ll;
typedef pair<int, int> P;
typedef vector<int> VI;
typedef vector<P> VP;

#define omajinai                                                               \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0)

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define RFOR(i, b, a) for (int i = (b)-1; i >= (a); --i)
#define RREP(i, n) RFOR(i, n, 0)

#define LFOR(i, a, b) for (ll i = (a); i < (b); ++i)
#define RLFOR(i, b, a) for (ll i = (b)-1; i >= (a); --i)

#define ALL(a) (a).begin(), (a).end()
#define UNIQUE(a) (a).erase(unique((a).begin(), (a).end()), (a).end())
#define MP make_pair
#define PB push_back
#define EACH(i, c) REP(i, (int)(c).size())
#define EXIST(s, e) ((s).find(e) != (s).end())
#define SORT(c) sort((c).begin(), (c).end())

#define dump(x) cerr << "[L " << __LINE__ << "] " << #x << " = " << (x) << "\\n";
#define dump2(x, y)                                                            \\
  cerr << "[L " << __LINE__ << "] " << #x << " = " << (x) << " , " << #y       \\
       << " = " << (y) << "\\n";

const int INF = 1e9;
const double EPS = 1e-10;
const int _N = 1e5;

int n, k, q;
VP edges[_N];

struct Node {
  ll top = 0, center = 0, bottom = 0;
  Node() {} //
  Node(ll top, ll center, ll bottom)
      : top(top), center(center), bottom(bottom) {}
};

const Node ZERO(-1, 0, 0);

Node merge(Node left, Node right) {
  if (left.top == -1)
    return right;
  if (right.top == -1)
    return left;

  ll center = left.center + right.center;
  ll len = left.bottom + right.top;
  if (len % k != 0)
    center += len; //

  return Node(left.top, center, right.bottom);
}

struct SegmentTree {
  Node ds[_N * 4 + 1];
  int m;
  Node temp;

  void init(int n) {
    m = 1;
    while (m < n)
      m *= 2;
    REP(i, m * 2 - 1)
    ds[i] = ZERO;
  }

  void update(int i, Node x) {
    i += m - 1;
    ds[i] = x;
    while (i >= 1) {
      i = (i - 1) / 2;
      ds[i] = merge(ds[i * 2 + 1], ds[i * 2 + 2]);
    }
  }

  Node get(int i) { return ds[i + m - 1]; }
  Node query(int a, int b) {
    temp = ZERO;
    _query(a, b, 0, 0, m);
    return temp;
  }

  void _query(int a, int b, int k, int l, int r) {
    if (r <= a || b <= l)
      return; //

    if (a <= l && r <= b) {
      temp = merge(temp, ds[k]);
      return;
    }
    _query(a, b, k * 2 + 1, l, (r + l) / 2);
    _query(a, b, k * 2 + 2, (r + l) / 2, r);
  }
};

struct X {};
X x;

int par[_N], head[_N], size[_N], depth[_N];
int indx[_N], id;

int dfs1(int t, int p) {
  par[t] = p;
  if (p != -1)
    depth[t] = depth[p] + 1; //
  int sz = 1;
  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    sz += dfs1(j, t);
  }
  return size[t] = sz;
}

SegmentTree seg;

void dfs2(int t, int p, int d) {
  //
  indx[t] = id;
  seg.update(id, Node(d, 0, 0));
  id++;

  int minP = -1;
  int minSz = INF;
  int minCost = 0;
  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    if (minSz > size[j]) {
      minP = j;
      minSz = size[j];
      minCost = edges[t][i].second;
    }
  }
  if (minP == -1)
    return;
  head[minP] = head[t];

  dfs2(minP, t, minCost); // for SegmentTree

  EACH(i, edges[t]) {
    int j = edges[t][i].first;
    if (j == p)
      continue;
    if (j == minP)
      continue;
    dfs2(j, t, edges[t][i].second);
  }
}

// HL-Decomposition
void decomposite() {
  dfs1(0, -1);
  REP(i, n) { head[i] = i; }
  dfs2(0, -1, 0);
}

int lca(int a, int b) {
  int p = head[a];
  int q = head[b];
  if (p == q)
    return depth[a] < depth[b] ? a : b;
  if (depth[p] < depth[q]) {
    return lca(a, par[q]);
  } else {
    return lca(par[p], b);
  }
}

ll dist(int p, int top) {
  Node d = ZERO;
  while (true) {
    if (head[p] == head[top]) {
      d = merge(seg.query(indx[top], indx[p] + 1), d);
      break;
    }
    d = merge(seg.query(indx[head[p]], indx[p] + 1), d);

    p = par[head[p]];
  }

  return d.center;
}

int main() {
  omajinai;
  cin >> n >> k;
  REP(i, n - 1) {
    int a, b, c;
    cin >> a >> b >> c;
    edges[a].PB(MP(b, c));
    edges[b].PB(MP(a, c));
  }

  seg.init(n);
  decomposite();

  cin >> q;

  REP(i, q) {
    string s;
    cin >> s;
    if (s[0] == 'a') { // add
      int x, d;
      cin >> x >> d;

      Node one = seg.get(indx[x]);
      one.top += d;
      one.bottom += d;
      seg.update(indx[x], one);
    } else { // send
      int s, t;
      cin >> s >> t;
      int l = lca(s, t);
      int d = dist(s, l) + dist(t, l);
      cout << d << endl;
    }
  }
}