Problem p02986 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <assert.h>
#include <complex>
#include <cstring>
#include <functional>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()
#define X first
#define Y second
#define pb(x) push_back(x)
#define FOR(k, n) for (int k = 0; k < n; k++)
#define FORN(k, s, n) for (int k = s; k < n; k++)
#define debug(x) printf("%s = %d\n", #x, x)
#define debug2(x, y) printf("%s = %d, %s = %d\n", #x, x, #y, y)
#define debug3(x, y, z)                                                        \
  printf("%s = %d, %s = %d, %s = %d\n", #x, x, #y, y, #z, z)
#define debug4(x, y, z, t)                                                     \
  printf("%s = %d, %s = %d, %s = %d, %s = %d\n", #x, x, #y, y, #z, z, #t, t)

#pragma warning(disable : 6031)
#pragma warning(disable : 4101)

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef long double ld;
typedef complex<double> cd;

const int INF = 1e9;
const long long LINF = 1e18;
const double PI = acos(-1);

// Solution

const int MAXN = 100005;
const int SQRT = 317;

struct Edge {
  ll c;
  ll d;

  Edge() : Edge(0, 0) {}
  Edge(ll c, ll d) : c{c}, d{d} {}
};

ll n, q, u, v, c, d, x, y, o_size, dtime, P[MAXN], S[MAXN], TIN[MAXN], D[MAXN],
    TOUT[MAXN], DP[MAXN][20];
vector<unordered_map<ll, Edge>> G;
vector<pair<ll, ll>> depths;
pair<ll, ll> M[2 * SQRT][MAXN];
map<ll, ll> CM;
char O[20 * MAXN];

void compute_root_info(ll v) {
  ll u = v;
  CM[v] = sz(CM);

  while (P[u] != -1) {
    M[CM[v]][G[u][P[u]].c].X++;
    M[CM[v]][G[u][P[u]].c].Y += G[u][P[u]].d;
    D[CM[v]] += G[u][P[u]].d;
    u = P[u];
  }
}

pair<ll, pair<ll, ll>> get_root_info(ll v, ll c) {
  ll tot = 0;
  pair<ll, ll> p;

  while (P[v] != -1) {
    if (S[v] == 1) {
      tot += D[CM[v]];
      p.X += M[CM[v]][c].X;
      p.Y += M[CM[v]][c].Y;

      return {tot, p};
    }

    tot += G[v][P[v]].d;

    if (G[v][P[v]].c == c) {
      p.X++;
      p.Y += G[v][P[v]].d;
    }

    v = P[v];
  }

  return {tot, p};
}

void dfs(ll v, ll p, ll depth) {
  P[v] = p;

  depths[v] = {depth, v};
  TIN[v] = dtime++;

  for (pair<ll, Edge> e : G[v]) {
    if (e.X != p) {
      dfs(e.X, v, depth + 1);
    }
  }

  TOUT[v] = dtime++;
}

bool is_ancestor(ll u, ll v) {
  return (TIN[u] <= TIN[v] && TOUT[v] <= TOUT[u]);
}

ll lca(ll u, ll v) {
  if (is_ancestor(u, v)) {
    return u;
  } else if (is_ancestor(v, u)) {
    return v;
  } else {
    for (int i = 19; i >= 0; i--) {
      if (!is_ancestor(DP[u][i], v)) {
        u = DP[u][i];
      }
    }

    return DP[u][0];
  }
}

void init_sparse_table() {
  for (int i = 0; i < n; i++) {
    DP[i][0] = P[i] == -1 ? i : P[i];
  }

  for (int i = 1; i < 20; i++) {
    for (int k = 0; k < n; k++) {
      DP[k][i] = DP[DP[k][i - 1]][i - 1];
    }
  }
}

int main() {
  scanf("%lld%lld", &n, &q);

  G.resize(n);
  depths.resize(n);

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

    G[u - 1][v - 1] = {c, d};
    G[v - 1][u - 1] = {c, d};
  }

  dfs(0, -1, 0);
  init_sparse_table();
  sort(depths.begin(), depths.end(), greater<pair<ll, ll>>());

  for (int i = 0; i < n; i++) {
    v = depths[i].Y;

    for (int k = 0; k < SQRT; k++) {
      if (S[v] == 1) {
        break;
      } else if (k == SQRT - 1 || P[v] == -1) {
        S[v] = 1;
        compute_root_info(v);
        break;
      }

      v = P[v];
    }
  }

  while (q--) {
    scanf("%lld%lld%lld%lld", &x, &y, &u, &v);

    pair<ll, pair<ll, ll>> u_stats = get_root_info(u - 1, x);
    pair<ll, pair<ll, ll>> v_stats = get_root_info(v - 1, x);
    pair<ll, pair<ll, ll>> lca_stats = get_root_info(lca(u - 1, v - 1), x);

    ll tot_dist = u_stats.X + v_stats.X - 2 * lca_stats.X;
    ll edge_cnt = u_stats.Y.X + v_stats.Y.X - 2 * lca_stats.Y.X;
    ll edge_weights = u_stats.Y.Y + v_stats.Y.Y - 2 * lca_stats.Y.Y;

    o_size +=
        sprintf(O + o_size, "%lld\n", tot_dist - edge_weights + y * edge_cnt);
  }

  printf("%s", O);

  return 0;
}

Fixed submission:
#include <algorithm>
#include <assert.h>
#include <complex>
#include <cstring>
#include <functional>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdio.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()
#define X first
#define Y second
#define pb(x) push_back(x)
#define FOR(k, n) for (int k = 0; k < n; k++)
#define FORN(k, s, n) for (int k = s; k < n; k++)
#define debug(x) printf("%s = %d\n", #x, x)
#define debug2(x, y) printf("%s = %d, %s = %d\n", #x, x, #y, y)
#define debug3(x, y, z)                                                        \
  printf("%s = %d, %s = %d, %s = %d\n", #x, x, #y, y, #z, z)
#define debug4(x, y, z, t)                                                     \
  printf("%s = %d, %s = %d, %s = %d, %s = %d\n", #x, x, #y, y, #z, z, #t, t)

#pragma warning(disable : 6031)
#pragma warning(disable : 4101)

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef long double ld;
typedef complex<double> cd;

const int INF = 1e9;
const long long LINF = 1e18;
const double PI = acos(-1);

// Solution

const int MAXN = 100005;
const int SQRT = 317;

struct Edge {
  ll c;
  ll d;

  Edge() : Edge(0, 0) {}
  Edge(ll c, ll d) : c{c}, d{d} {}
};

ll n, q, u, v, c, d, x, y, o_size, dtime, P[MAXN], S[MAXN], TIN[MAXN], D[MAXN],
    TOUT[MAXN], DP[MAXN][20];
vector<map<ll, Edge>> G;
vector<pair<ll, ll>> depths;
pair<ll, ll> M[2 * SQRT][MAXN];
map<ll, ll> CM;
char O[20 * MAXN];

void compute_root_info(ll v) {
  ll u = v;
  CM[v] = sz(CM);

  while (P[u] != -1) {
    M[CM[v]][G[u][P[u]].c].X++;
    M[CM[v]][G[u][P[u]].c].Y += G[u][P[u]].d;
    D[CM[v]] += G[u][P[u]].d;
    u = P[u];
  }
}

pair<ll, pair<ll, ll>> get_root_info(ll v, ll c) {
  ll tot = 0;
  pair<ll, ll> p;

  while (P[v] != -1) {
    if (S[v] == 1) {
      tot += D[CM[v]];
      p.X += M[CM[v]][c].X;
      p.Y += M[CM[v]][c].Y;

      return {tot, p};
    }

    tot += G[v][P[v]].d;

    if (G[v][P[v]].c == c) {
      p.X++;
      p.Y += G[v][P[v]].d;
    }

    v = P[v];
  }

  return {tot, p};
}

void dfs(ll v, ll p, ll depth) {
  P[v] = p;

  depths[v] = {depth, v};
  TIN[v] = dtime++;

  for (pair<ll, Edge> e : G[v]) {
    if (e.X != p) {
      dfs(e.X, v, depth + 1);
    }
  }

  TOUT[v] = dtime++;
}

bool is_ancestor(ll u, ll v) {
  return (TIN[u] <= TIN[v] && TOUT[v] <= TOUT[u]);
}

ll lca(ll u, ll v) {
  if (is_ancestor(u, v)) {
    return u;
  } else if (is_ancestor(v, u)) {
    return v;
  } else {
    for (int i = 19; i >= 0; i--) {
      if (!is_ancestor(DP[u][i], v)) {
        u = DP[u][i];
      }
    }

    return DP[u][0];
  }
}

void init_sparse_table() {
  for (int i = 0; i < n; i++) {
    DP[i][0] = P[i] == -1 ? i : P[i];
  }

  for (int i = 1; i < 20; i++) {
    for (int k = 0; k < n; k++) {
      DP[k][i] = DP[DP[k][i - 1]][i - 1];
    }
  }
}

int main() {
  scanf("%lld%lld", &n, &q);

  G.resize(n);
  depths.resize(n);

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

    G[u - 1][v - 1] = {c, d};
    G[v - 1][u - 1] = {c, d};
  }

  dfs(0, -1, 0);
  init_sparse_table();
  sort(depths.begin(), depths.end(), greater<pair<ll, ll>>());

  for (int i = 0; i < n; i++) {
    v = depths[i].Y;

    for (int k = 0; k < SQRT; k++) {
      if (S[v] == 1) {
        break;
      } else if (k == SQRT - 1 || P[v] == -1) {
        S[v] = 1;
        compute_root_info(v);
        break;
      }

      v = P[v];
    }
  }

  while (q--) {
    scanf("%lld%lld%lld%lld", &x, &y, &u, &v);

    pair<ll, pair<ll, ll>> u_stats = get_root_info(u - 1, x);
    pair<ll, pair<ll, ll>> v_stats = get_root_info(v - 1, x);
    pair<ll, pair<ll, ll>> lca_stats = get_root_info(lca(u - 1, v - 1), x);

    ll tot_dist = u_stats.X + v_stats.X - 2 * lca_stats.X;
    ll edge_cnt = u_stats.Y.X + v_stats.Y.X - 2 * lca_stats.Y.X;
    ll edge_weights = u_stats.Y.Y + v_stats.Y.Y - 2 * lca_stats.Y.Y;

    o_size +=
        sprintf(O + o_size, "%lld\n", tot_dist - edge_weights + y * edge_cnt);
  }

  printf("%s", O);

  return 0;
}