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

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

#define CIN_ONLY if (1)
struct cww {
  cww() {
    CIN_ONLY {
      ios::sync_with_stdio(false);
      cin.tie(0);
    }
  }
} star;
#define fin "\\n"
#define FOR(i, bg, ed) for (int i = (bg); i < (ed); i++)
#define REP(i, n) FOR(i, 0, n)
#define ALL(v) (v).begin(), (v).end()
#define fi first
#define se second
#define pb push_back
#define DEBUG if (0)
#define REC(ret, ...) std::function<ret(__VA_ARGS__)>
template <typename T> inline bool chmin(T &l, T r) {
  bool a = l > r;
  if (a)
    l = r;
  return a;
}
template <typename T> inline bool chmax(T &l, T r) {
  bool a = l < r;
  if (a)
    l = r;
  return a;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &it : v)
    is >> it;
  return is;
}
typedef vector<int> V;
typedef vector<V> VV;
typedef VV Graph;
typedef pair<int, int> P;

int pp[22][312345];
int s[612345];

void euler_tour(int v, Graph &C, vector<P> &seg, int &cnt) {
  int pt = 0;
  s[pt++] = -1;
  s[pt++] = 1;
  while (pt) {
    int vv = s[--pt];
    int v = abs(vv) - 1;
    if (vv > 0) {
      seg[v].fi = cnt++;
      for (auto &u : C[v])
        if (u != pp[0][v]) {
          s[pt++] = -u - 1;
          s[pt++] = u + 1;
        }
    } else
      seg[v].se = cnt++;
  }
}

// lca?????????
void build_PP(int N) {
  for (int k = 0, f = 1; f; k++) {
    f = 0;
    for (int i = 0; i < N; i++) {
      if (pp[k][i] < 0)
        pp[k + 1][i] = -1;
      else {
        pp[k + 1][i] = pp[k][pp[k][i]];
        f = 1;
      }
    }
  }
}

int lca(int u, int v, V &D) {
  if (D[u] > D[v])
    swap(u, v);
  for (int k = 0, d; (d = D[v] - D[u]); k++)
    if ((d >> k) & 1)
      v = pp[k][v];
  if (u == v)
    return v;
  for (int k = 21; k >= 0; k--) {
    if (pp[k][u] != pp[k][v]) {
      u = pp[k][u];
      v = pp[k][v];
    }
  }
  return pp[0][u];
}
LL depth[312345];
LL query[312345];
void dfs(int v, Graph &g, V &d) {
  int pt = 0;
  s[pt++] = v;
  while (pt) {
    int v = s[--pt];
    depth[v] = query[v] = 0;
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        d[u] = d[v] + 1;
        pp[0][u] = v;
        s[pt++] = u;
      }
  }
}
void querydfs(int vv, Graph &g) {
  int pt = 0;
  s[pt++] = vv;
  while (pt) {
    int v = s[--pt];
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        query[u] += query[v];
        s[pt++] = u;
      }
  }
  s[pt++] = vv;
  while (pt) {
    int v = s[--pt];
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        query[u] += query[v];
        depth[u] += query[v];
        s[pt++] = u;
      }
    query[v] = 0;
  }
}

const int B = 1024;
int main() {
  int N, Q;
  cin >> N >> Q;
  Graph g(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    g[a].pb(b);
    g[b].pb(a);
  }
  V d(N, 0);
  REP(i, 22) REP(j, N) pp[i][j] = -1;
  dfs(0, g, d);
  build_PP(N);
  vector<P> seg(N);
  vector<LL> qx, qy;
  qx.reserve(B);
  qy.reserve(B);
  int M = 1;
  euler_tour(0, g, seg, M);
  REP(q, Q) {
    int type, x, y;
    cin >> type >> x >> y;
    if (type == 0) {
      int z = lca(x, y, d);
      LL A = depth[x];
      LL B = depth[y];
      LL C = depth[z];
      REP(i, qx.size()) {
        int l = seg[qx[i]].fi;
        int r = seg[qx[i]].se;

        if (l < seg[x].fi && seg[x].se < r) {
          LL dist = d[x] - d[qx[i]];
          A += qy[i] * dist;
        }
        if (l < seg[y].fi && seg[y].se < r) {
          LL dist = d[y] - d[qx[i]];
          B += qy[i] * dist;
        }
        if (l < seg[z].fi && seg[z].se < r) {
          LL dist = d[z] - d[qx[i]];
          C += qy[i] * dist;
        }
      }
      // cout<<A<<endl;
      cout << A + B - 2 * C << endl;
    } else {
      qx.pb(x);
      qy.pb(y);
    }
    if (qx.size() == B) {
      REP(i, B) { query[qx[i]] += qy[i]; }
      querydfs(0, g);
      qx.clear();
      qy.clear();
    }
  }
  return 0;
}

Error: TLE

Fixed submission:

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

#define CIN_ONLY if (1)
struct cww {
  cww() {
    CIN_ONLY {
      ios::sync_with_stdio(false);
      cin.tie(0);
    }
  }
} star;
#define fin "\\n"
#define FOR(i, bg, ed) for (int i = (bg); i < (ed); i++)
#define REP(i, n) FOR(i, 0, n)
#define ALL(v) (v).begin(), (v).end()
#define fi first
#define se second
#define pb push_back
#define DEBUG if (0)
#define REC(ret, ...) std::function<ret(__VA_ARGS__)>
template <typename T> inline bool chmin(T &l, T r) {
  bool a = l > r;
  if (a)
    l = r;
  return a;
}
template <typename T> inline bool chmax(T &l, T r) {
  bool a = l < r;
  if (a)
    l = r;
  return a;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &it : v)
    is >> it;
  return is;
}
typedef vector<int> V;
typedef vector<V> VV;
typedef VV Graph;
typedef pair<int, int> P;

int pp[22][312345];
int s[612345];

void euler_tour(int v, Graph &C, vector<P> &seg, int &cnt) {
  int pt = 0;
  s[pt++] = -1;
  s[pt++] = 1;
  while (pt) {
    int vv = s[--pt];
    int v = abs(vv) - 1;
    if (vv > 0) {
      seg[v].fi = cnt++;
      for (auto &u : C[v])
        if (u != pp[0][v]) {
          s[pt++] = -u - 1;
          s[pt++] = u + 1;
        }
    } else
      seg[v].se = cnt++;
  }
}

// lca?????????
void build_PP(int N) {
  for (int k = 0, f = 1; f; k++) {
    f = 0;
    for (int i = 0; i < N; i++) {
      if (pp[k][i] < 0)
        pp[k + 1][i] = -1;
      else {
        pp[k + 1][i] = pp[k][pp[k][i]];
        f = 1;
      }
    }
  }
}

int lca(int u, int v, V &D) {
  if (D[u] > D[v])
    swap(u, v);
  for (int k = 0, d; (d = D[v] - D[u]); k++)
    if ((d >> k) & 1)
      v = pp[k][v];
  if (u == v)
    return v;
  for (int k = 21; k >= 0; k--) {
    if (pp[k][u] != pp[k][v]) {
      u = pp[k][u];
      v = pp[k][v];
    }
  }
  return pp[0][u];
}
LL depth[312345];
LL query[312345];
void dfs(int v, Graph &g, V &d) {
  int pt = 0;
  s[pt++] = v;
  while (pt) {
    int v = s[--pt];
    depth[v] = query[v] = 0;
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        d[u] = d[v] + 1;
        pp[0][u] = v;
        s[pt++] = u;
      }
  }
}
void querydfs(int vv, Graph &g) {
  int pt = 0;
  s[pt++] = vv;
  while (pt) {
    int v = s[--pt];
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        query[u] += query[v];
        s[pt++] = u;
      }
  }
  s[pt++] = vv;
  while (pt) {
    int v = s[--pt];
    for (auto &u : g[v])
      if (pp[0][v] != u) {
        query[u] += query[v];
        depth[u] += query[v];
        s[pt++] = u;
      }
    query[v] = 0;
  }
}

int main() {
  int N, Q;
  cin >> N >> Q;
  Graph g(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    g[a].pb(b);
    g[b].pb(a);
  }
  V d(N, 0);
  REP(i, 22) REP(j, N) pp[i][j] = -1;
  dfs(0, g, d);
  build_PP(N);
  vector<P> seg(N);
  vector<LL> qx, qy;
  qx.reserve(512);
  qy.reserve(512);
  int M = 1;
  euler_tour(0, g, seg, M);
  REP(q, Q) {
    int type, x, y;
    cin >> type >> x >> y;
    if (type == 0) {
      int z = lca(x, y, d);
      LL A = depth[x];
      LL B = depth[y];
      LL C = depth[z];
      REP(i, qx.size()) {
        int l = seg[qx[i]].fi;
        int r = seg[qx[i]].se;

        if (l < seg[x].fi && seg[x].se < r) {
          LL dist = d[x] - d[qx[i]];
          A += qy[i] * dist;
        }
        if (l < seg[y].fi && seg[y].se < r) {
          LL dist = d[y] - d[qx[i]];
          B += qy[i] * dist;
        }
        if (l < seg[z].fi && seg[z].se < r) {
          LL dist = d[z] - d[qx[i]];
          C += qy[i] * dist;
        }
      }
      // cout<<A<<endl;
      cout << A + B - 2 * C << endl;
    } else {
      qx.pb(x);
      qy.pb(y);
    }
    if (qx.size() == 512) {
      REP(i, 512) { query[qx[i]] += qy[i]; }
      querydfs(0, g);
      qx.clear();
      qy.clear();
    }
  }
  return 0;
}