Problem p02986 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

ostream &operator<<(ostream &out, string str) {
  for (char c : str)
    cout << c;
  return out;
}

template <class L, class R> ostream &operator<<(ostream &out, pair<L, R> p) {
  return out << "(" << p.first << ", " << p.second << ")";
}

template <class T>
auto operator<<(ostream &out, T a) -> decltype(a.begin(), out) {
  out << "{";
  for (auto it = a.begin(); it != a.end(); it = next(it))
    out << (it != a.begin() ? ", " : "") << *it;
  return out << "}";
}

void dump() { cerr << "\n"; }
template <class T, class... Ts> void dump(T a, Ts... x) {
  cerr << a << ", ";
  dump(x...);
}

#ifdef DEBUG
#define debug(...) cerr << "[" #__VA_ARGS__ "]: ", dump(__VA_ARGS__)
#else
#define debug(...) false
#endif

#define REP(i, n) for (int i = 0; i < n; i++)
#define FOR(i, a, b) for (int i = a; i <= b; i++)
#define ST first
#define ND second

template <class T> int size(T &&a) { return a.size(); }

using LL = long long;
using PII = pair<int, int>;

struct edge {
  int to, col, len;
};

vector<vector<edge>> graph;
vector<int> par, dist;
vector<int> pre, pos;

int t = 0;
void set_par(int v = 0, int p = -1) {
  pre[v] = t++;
  par[v] = p;
  for (edge &e : graph[v]) {
    int u = e.to;
    int d = e.len;
    if (u != p) {
      dist[u] = dist[v] + d;
      set_par(u, v);
    }
  }
  pos[v] = t++;
}

bool is_par(int v, int u) { return pre[v] <= pre[u] && pos[u] <= pos[v]; }

vector<vector<int>> jump;
int LCA(int v, int u) {
  if (is_par(v, u))
    return v;
  if (is_par(u, v))
    return u;

  for (int i = 19; i >= 0; i--) {
    if (jump[i][v] != -1 && !is_par(jump[i][v], u))
      v = jump[i][v];
  }

  return jump[0][v];
}

vector<vector<PII>> dists;
void dist_dfs(int c, int v = 0) {
  for (edge &e : graph[v]) {
    int u = e.to;
    int d = e.len;
    int col = e.col;
    if (u != par[v]) {
      dists[c][u] = dists[c][v];
      if (col == c)
        dists[c][u].ND++;
      else
        dists[c][u].ST += d;
      dist_dfs(c, u);
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int n, q;
  cin >> n >> q;

  struct just_edge {
    int a, b, d;
  };

  vector<vector<just_edge>> edges(n - 1);
  vector<int> color_cnt(n - 1);
  graph.resize(n);
  REP(i, n - 1) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    a--, b--, c--;
    color_cnt[c]++;
    graph[a].emplace_back(edge{b, c, d});
    graph[b].emplace_back(edge{a, c, d});
    edges[c].emplace_back(just_edge{a, b, d});
  }

  par.resize(n);
  dist.resize(n);
  pre.resize(n);
  pos.resize(n);
  set_par();

  int log = 20;
  jump.resize(20, vector<int>(n));
  jump[0] = par;

  FOR(i, 1, log - 1) {
    REP(j, n)
    jump[i][j] = (jump[i - 1][j] == -1 ? -1 : jump[i - 1][jump[i - 1][j]]);
  }

  int S = 0;
  dists.resize(n - 1);
  REP(i, n - 1) {
    if (color_cnt[i] >= S) {
      dists[i].resize(n);
      dist_dfs(i);
    }
  }

  auto add = [&](PII a, PII b) -> PII { return {a.ST + b.ST, a.ND + b.ND}; };

  auto rem = [&](PII a, PII b) -> PII { return {a.ST - b.ST, a.ND - b.ND}; };

  REP(i, q) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    x--, u--, v--;

    if (color_cnt[x] >= S) {
      int lca = LCA(u, v);
      PII ret =
          add(rem(dists[x][u], dists[x][lca]), rem(dists[x][v], dists[x][lca]));
      cout << ret.ST + ret.ND * y << "\n";
    } else {
      int lca = LCA(u, v);
      int ret = dist[u] + dist[v] - dist[lca] * 2;
      for (auto &e : edges[x]) {
        int a = e.a;
        int b = e.b;
        int d = e.d;
        if (is_par(a, b))
          swap(a, b);
        if (!is_par(a, lca) && (is_par(a, u) || is_par(a, v)))
          ret += y - d;
      }
      cout << ret << "\n";
    }
  }
}


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

ostream &operator<<(ostream &out, string str) {
  for (char c : str)
    cout << c;
  return out;
}

template <class L, class R> ostream &operator<<(ostream &out, pair<L, R> p) {
  return out << "(" << p.first << ", " << p.second << ")";
}

template <class T>
auto operator<<(ostream &out, T a) -> decltype(a.begin(), out) {
  out << "{";
  for (auto it = a.begin(); it != a.end(); it = next(it))
    out << (it != a.begin() ? ", " : "") << *it;
  return out << "}";
}

void dump() { cerr << "\n"; }
template <class T, class... Ts> void dump(T a, Ts... x) {
  cerr << a << ", ";
  dump(x...);
}

#ifdef DEBUG
#define debug(...) cerr << "[" #__VA_ARGS__ "]: ", dump(__VA_ARGS__)
#else
#define debug(...) false
#endif

#define REP(i, n) for (int i = 0; i < n; i++)
#define FOR(i, a, b) for (int i = a; i <= b; i++)
#define ST first
#define ND second

template <class T> int size(T &&a) { return a.size(); }

using LL = long long;
using PII = pair<int, int>;

struct edge {
  int to, col, len;
};

vector<vector<edge>> graph;
vector<int> par, dist;
vector<int> pre, pos;

int t = 0;
void set_par(int v = 0, int p = -1) {
  pre[v] = t++;
  par[v] = p;
  for (edge &e : graph[v]) {
    int u = e.to;
    int d = e.len;
    if (u != p) {
      dist[u] = dist[v] + d;
      set_par(u, v);
    }
  }
  pos[v] = t++;
}

bool is_par(int v, int u) { return pre[v] <= pre[u] && pos[u] <= pos[v]; }

vector<vector<int>> jump;
int LCA(int v, int u) {
  if (is_par(v, u))
    return v;
  if (is_par(u, v))
    return u;

  for (int i = 19; i >= 0; i--) {
    if (jump[i][v] != -1 && !is_par(jump[i][v], u))
      v = jump[i][v];
  }

  return jump[0][v];
}

vector<vector<PII>> dists;
void dist_dfs(int c, int v = 0) {
  for (edge &e : graph[v]) {
    int u = e.to;
    int d = e.len;
    int col = e.col;
    if (u != par[v]) {
      dists[c][u] = dists[c][v];
      if (col == c)
        dists[c][u].ND++;
      else
        dists[c][u].ST += d;
      dist_dfs(c, u);
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int n, q;
  cin >> n >> q;

  struct just_edge {
    int a, b, d;
  };

  vector<vector<just_edge>> edges(n - 1);
  vector<int> color_cnt(n - 1);
  graph.resize(n);
  REP(i, n - 1) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    a--, b--, c--;
    color_cnt[c]++;
    graph[a].emplace_back(edge{b, c, d});
    graph[b].emplace_back(edge{a, c, d});
    edges[c].emplace_back(just_edge{a, b, d});
  }

  par.resize(n);
  dist.resize(n);
  pre.resize(n);
  pos.resize(n);
  set_par();

  int log = 20;
  jump.resize(20, vector<int>(n));
  jump[0] = par;

  FOR(i, 1, log - 1) {
    REP(j, n)
    jump[i][j] = (jump[i - 1][j] == -1 ? -1 : jump[i - 1][jump[i - 1][j]]);
  }

  int S = 320;
  dists.resize(n - 1);
  REP(i, n - 1) {
    if (color_cnt[i] >= S) {
      dists[i].resize(n);
      dist_dfs(i);
    }
  }

  auto add = [&](PII a, PII b) -> PII { return {a.ST + b.ST, a.ND + b.ND}; };

  auto rem = [&](PII a, PII b) -> PII { return {a.ST - b.ST, a.ND - b.ND}; };

  REP(i, q) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    x--, u--, v--;

    if (color_cnt[x] >= S) {
      int lca = LCA(u, v);
      PII ret =
          add(rem(dists[x][u], dists[x][lca]), rem(dists[x][v], dists[x][lca]));
      cout << ret.ST + ret.ND * y << "\n";
    } else {
      int lca = LCA(u, v);
      int ret = dist[u] + dist[v] - dist[lca] * 2;
      for (auto &e : edges[x]) {
        int a = e.a;
        int b = e.b;
        int d = e.d;
        if (is_par(a, b))
          swap(a, b);
        if (!is_par(a, lca) && (is_par(a, u) || is_par(a, v)))
          ret += y - d;
      }
      cout << ret << "\n";
    }
  }
}