Problem p02986 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;
typedef complex<ld> cd;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<ld, ld> pd;

typedef vector<int> vi;
typedef vector<ld> vd;
typedef vector<ll> vl;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<cd> vcd;

template <class T>
using Tree =
    tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

#define FOR(i, a, b) for (int i = a; i < (b); i++)
#define F0R(i, a) for (int i = 0; i < (a); i++)
#define FORd(i, a, b) for (int i = (b)-1; i >= a; i--)
#define F0Rd(i, a) for (int i = (a)-1; i >= 0; i--)

#define sz(x) (int)(x).size()
#define mp make_pair
#define pb push_back
#define f first
#define s second
#define lb lower_bound
#define ub upper_bound
#define all(x) x.begin(), x.end()

const int MOD = 1000000007;
const ll INF = 1e18;
const int MX = 100001; // check the limits, dummy

vector<vpi> graph;
int level[MX];
int parent[MX];
bool special[MX];
bool checked[MX];
vpi edgeData;
vi parentEdge;
vector<map<int, pi>> colorData;
map<int, pi> curColors;
int lcaData[MX][18];
int dist[MX];

void dfs1(int v, int p) {
  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    level[nxt] = level[v] + 1;
    parent[nxt] = v;
    parentEdge[nxt] = graph[v][i].s;
    dist[nxt] = dist[v] + edgeData[parentEdge[nxt]].s;
    dfs1(nxt, v);
  }
}

void dfs2(int v, int p) {
  if (special[v]) {
    colorData[v] = curColors;
  }

  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    int ind = graph[v][i].s;
    int col = edgeData[ind].f;
    int dist = edgeData[ind].s;
    if (!curColors.count(col)) {
      curColors.insert(mp(col, mp(0, 0)));
    }
    curColors[col].f++;
    curColors[col].s += dist;
    dfs2(nxt, v);
    curColors[col].f--;
    curColors[col].s -= dist;
    if (curColors[col].f == 0) {
      curColors.erase(col);
    }
  }
}

void dfs3(int v, int p) {
  lcaData[v][0] = p;
  FOR(i, 1, 18) {
    if (lcaData[v][i - 1] != -1) {
      lcaData[v][i] = lcaData[lcaData[v][i - 1]][i - 1];
    }
  }

  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    dfs3(nxt, v);
  }
}

int lca(int u, int v) {
  if (level[u] < level[v]) {
    swap(u, v);
  }

  F0Rd(i, 18) {
    if (level[u] - (1 << i) >= level[v]) {
      u = lcaData[u][i];
    }
  }

  if (u == v)
    return u;

  F0Rd(i, 18) {
    if (lcaData[u][i] != -1 && lcaData[v][i] != -1 &&
        lcaData[u][i] != lcaData[v][i]) {
      u = lcaData[u][i];
      v = lcaData[v][i];
    }
  }

  return lcaData[u][0];
}

pi getData(int v, int color) {
  if (special[v]) {
    if (!colorData[v].count(color)) {
      return mp(0, 0);
    } else {
      return colorData[v][color];
    }
  }

  pi prev = getData(parent[v], color);
  int ind = parentEdge[v];
  if (edgeData[ind].f == color) {
    prev.f++;
    prev.s += edgeData[ind].s;
  }
  return prev;
}

int rootDist(int v, pi data, int newLen) {
  return dist[v] - data.s + data.f * newLen;
}

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

  int N;
  cin >> N;
  int Q;
  cin >> Q;
  F0R(i, N) {
    vpi blank;
    graph.pb(blank);
    parentEdge.pb(0);
  }
  F0R(i, N - 1) {
    int A, B, C, D;
    cin >> A >> B >> C >> D;
    A--;
    B--;
    edgeData.pb(mp(C, D));
    graph[A].pb(mp(B, i));
    graph[B].pb(mp(A, i));
  }

  parent[0] = -1;
  level[0] = 0;
  dist[0] = 0;
  dfs1(0, -1);

  priority_queue<pi> pq;
  F0R(i, N) { pq.push(mp(level[i], i)); }
  F0R(i, N) {
    special[i] = false;
    checked[i] = false;
  }
  special[0] = true;
  checked[0] = true;

  while (!pq.empty()) {
    int cur = pq.top().s;
    pq.pop();
    if (!pq.empty())
      assert(level[cur] <= pq.top().f);
    if (checked[cur])
      continue;
    F0R(i, 350) {
      cur = parent[cur];
      if (special[cur])
        goto done;
      checked[cur] = true;
    }
    special[cur] = true;
  done:;
  }

  F0R(i, N) {
    map<int, pi> blank;
    colorData.pb(blank);
  }

  dfs2(0, -1);

  F0R(i, N) {
    F0R(j, 18) { lcaData[i][j] = -1; }
  }

  dfs3(0, -1);

  F0R(query, Q) {
    int A, B, C, D;
    cin >> A >> B >> C >> D;
    C--;
    D--;
    int E = lca(C, D);
    cout << rootDist(C, getData(C, A), B) + rootDist(D, getData(D, A), B) -
                2 * rootDist(E, getData(E, A), B)
         << endl;
  }

  return 0;
}

// read the question correctly (ll vs int)
// template by bqi343


Fixed submission:
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;
typedef complex<ld> cd;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<ld, ld> pd;

typedef vector<int> vi;
typedef vector<ld> vd;
typedef vector<ll> vl;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<cd> vcd;

template <class T>
using Tree =
    tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

#define FOR(i, a, b) for (int i = a; i < (b); i++)
#define F0R(i, a) for (int i = 0; i < (a); i++)
#define FORd(i, a, b) for (int i = (b)-1; i >= a; i--)
#define F0Rd(i, a) for (int i = (a)-1; i >= 0; i--)

#define sz(x) (int)(x).size()
#define mp make_pair
#define pb push_back
#define f first
#define s second
#define lb lower_bound
#define ub upper_bound
#define all(x) x.begin(), x.end()

const int MOD = 1000000007;
const ll INF = 1e18;
const int MX = 100001; // check the limits, dummy

vector<vpi> graph;
int level[MX];
int parent[MX];
bool special[MX];
bool checked[MX];
vpi edgeData;
vi parentEdge;
vector<map<int, pi>> colorData;
map<int, pi> curColors;
int lcaData[MX][18];
int dist[MX];

void dfs1(int v, int p) {
  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    level[nxt] = level[v] + 1;
    parent[nxt] = v;
    parentEdge[nxt] = graph[v][i].s;
    dist[nxt] = dist[v] + edgeData[parentEdge[nxt]].s;
    dfs1(nxt, v);
  }
}

void dfs2(int v, int p) {
  if (special[v]) {
    colorData[v] = curColors;
  }

  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    int ind = graph[v][i].s;
    int col = edgeData[ind].f;
    int dist = edgeData[ind].s;
    if (!curColors.count(col)) {
      curColors.insert(mp(col, mp(0, 0)));
    }
    curColors[col].f++;
    curColors[col].s += dist;
    dfs2(nxt, v);
    curColors[col].f--;
    curColors[col].s -= dist;
    if (curColors[col].f == 0) {
      curColors.erase(col);
    }
  }
}

void dfs3(int v, int p) {
  lcaData[v][0] = p;
  FOR(i, 1, 18) {
    if (lcaData[v][i - 1] != -1) {
      lcaData[v][i] = lcaData[lcaData[v][i - 1]][i - 1];
    }
  }

  F0R(i, sz(graph[v])) {
    int nxt = graph[v][i].f;
    if (nxt == p)
      continue;
    dfs3(nxt, v);
  }
}

int lca(int u, int v) {
  if (level[u] < level[v]) {
    swap(u, v);
  }

  F0Rd(i, 18) {
    if (level[u] - (1 << i) >= level[v]) {
      u = lcaData[u][i];
    }
  }

  if (u == v)
    return u;

  F0Rd(i, 18) {
    if (lcaData[u][i] != -1 && lcaData[v][i] != -1 &&
        lcaData[u][i] != lcaData[v][i]) {
      u = lcaData[u][i];
      v = lcaData[v][i];
    }
  }

  return lcaData[u][0];
}

pi getData(int v, int color) {
  if (special[v]) {
    if (!colorData[v].count(color)) {
      return mp(0, 0);
    } else {
      return colorData[v][color];
    }
  }

  pi prev = getData(parent[v], color);
  int ind = parentEdge[v];
  if (edgeData[ind].f == color) {
    prev.f++;
    prev.s += edgeData[ind].s;
  }
  return prev;
}

int rootDist(int v, pi data, int newLen) {
  return dist[v] - data.s + data.f * newLen;
}

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

  int N;
  cin >> N;
  int Q;
  cin >> Q;
  F0R(i, N) {
    vpi blank;
    graph.pb(blank);
    parentEdge.pb(0);
  }
  F0R(i, N - 1) {
    int A, B, C, D;
    cin >> A >> B >> C >> D;
    A--;
    B--;
    edgeData.pb(mp(C, D));
    graph[A].pb(mp(B, i));
    graph[B].pb(mp(A, i));
  }

  parent[0] = -1;
  level[0] = 0;
  dist[0] = 0;
  dfs1(0, -1);

  priority_queue<pi> pq;
  F0R(i, N) { pq.push(mp(level[i], i)); }
  F0R(i, N) {
    special[i] = false;
    checked[i] = false;
  }
  special[0] = true;
  checked[0] = true;

  while (!pq.empty()) {
    int cur = pq.top().s;
    pq.pop();
    if (!pq.empty())
      assert(level[cur] >= pq.top().f);
    if (checked[cur])
      continue;
    F0R(i, 350) {
      cur = parent[cur];
      if (special[cur])
        goto done;
      checked[cur] = true;
    }
    special[cur] = true;
  done:;
  }

  F0R(i, N) {
    map<int, pi> blank;
    colorData.pb(blank);
  }

  dfs2(0, -1);

  F0R(i, N) {
    F0R(j, 18) { lcaData[i][j] = -1; }
  }

  dfs3(0, -1);

  F0R(query, Q) {
    int A, B, C, D;
    cin >> A >> B >> C >> D;
    C--;
    D--;
    int E = lca(C, D);
    cout << rootDist(C, getData(C, A), B) + rootDist(D, getData(D, A), B) -
                2 * rootDist(E, getData(E, A), B)
         << endl;
  }

  return 0;
}

// read the question correctly (ll vs int)
// template by bqi343