Problem p02986 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
#include <vector>
using namespace std;
const int MAXN = 1e5 + 5;
map<int, pair<int, int>> curColours;
vector<map<int, pair<int, int>>> colourData;

vector<pair<int, int>> v1[MAXN];
int level[MAXN];
int parent[MAXN];
int parentedge[MAXN];
int dist[MAXN];
vector<pair<int, int>> edgedata;
int dp[MAXN][24];
bool special[MAXN];
bool checkedi[MAXN];
void dfs1(int curr, int par) {
  for (auto z : v1[curr]) {
    int x = z.first;
    if (x == par) {
      continue;
    }
    level[x] = level[curr] + 1;
    parent[x] = curr;
    parentedge[x] = z.second;
    dist[x] = dist[curr] + edgedata[parentedge[x]].second;
    dfs1(x, curr);
  }
}
void dfs2(int curr, int par) {

  if (special[curr]) {

    colourData[curr] = curColours;
  }
  for (auto x : v1[curr]) {

    if (x.first == par) {
      continue;
    }
    int ind = x.second;
    // cout<<curr<<" "<<x.first<<" "<<ind<<endl;
    int col = edgedata[ind].first;
    int dist = edgedata[ind].second;
    if (!curColours.count(col)) {
      curColours.insert(make_pair(col, make_pair(0, 0)));
    }
    curColours[col].first++;
    curColours[col].second += dist;
    dfs2(x.first, curr);

    curColours[col].first--;
    curColours[col].second -= dist;

    if (curColours[col].first == 0) {
      curColours.erase(col);
    }
  }
}
int lca(int u, int v) {
  if (level[u] < level[v]) {
    swap(u, v);
  }
  for (int i = 20; i >= 0; i--) {
    if (level[u] - (1 << i) >= level[v]) {
      u = dp[u][i];
    }
  }
  if (u == v) {
    return u;
  }
  for (int i = 20; i >= 0; i--) {
    if (dp[u][i] != dp[v][i] && dp[u][i] != -1 && dp[v][i] != -1) {
      u = dp[u][i];
      v = dp[v][i];
    }
  }
  return parent[u];
}
pair<int, int> getdata(int v, int colour) {
  if (v == -1) {
    cout << 1234567 << endl;
    return make_pair(0, 0);
  }
  //   cout<<v<<" "<<parent[v]<<" "<<special[v]<<endl;
  if (special[v]) {
    if (!colourData[v].count(colour)) {
      return make_pair(0, 0);
    } else {
      return colourData[v][colour];
    }
  }

  pair<int, int> prev = getdata(parent[v], colour);
  int ind = parentedge[v];
  if (edgedata[ind].first == colour) {
    prev.first++;
    prev.second += edgedata[ind].second;
  }
  return prev;
}
int rootdist(int curr, pair<int, int> data, int newLen) {
  return dist[curr] - data.second + (data.first * newLen);
}
int main() {
  int n;
  cin >> n;
  int q;
  cin >> q;
  for (int i = 0; i < n - 1; i++) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    edgedata.push_back(make_pair(c, d));
    v1[a].push_back(make_pair(b, i));
    v1[b].push_back(make_pair(a, i));
  }
  parent[1] = -1;
  level[1] = 0;
  dist[1] = 0;
  dfs1(1, -1);
  priority_queue<pair<int, int>> pq;
  for (int i = 1; i <= n; i++) {
    pq.push(make_pair(level[i], i));
  }
  for (int i = 0; i <= n; i++) {
    special[i] = false;
    checkedi[i] = false;
  }

  int cnt = 1;
  special[1] = true;
  checkedi[1] = true;
  while (!pq.empty()) {
    int curr = pq.top().second;
    pq.pop();

    if (checkedi[curr]) {
      continue;
    }
    for (int i = 0; i < 350; i++) {
      curr = parent[curr];
      if (special[curr]) {
        break;
      }
      checkedi[curr] = true;
    }
    special[curr] = true;
  }

  for (int i = 0; i <= n; i++) {
    map<int, pair<int, int>> blank;
    colourData.push_back(blank);
  }
  dfs2(1, -1);
  for (int i = 1; i <= n; i++) {
    for (int j = 0; j <= 20; j++) {
      dp[i][j] = -1;
    }
  }
  for (int i = 1; i <= n; i++) {
    dp[i][0] = parent[i];
  }
  for (int j = 1; j <= n; j++) {
    for (int i = 1; i <= n; i++) {
      if (dp[i][j - 1] == -1) {
        dp[i][j] = -1;
      } else {
        dp[i][j] = max(dp[i][j], dp[dp[i][j - 1]][j - 1]);
      }
    }
  }
  while (q--) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    int e = lca(c, d);
    // cout<<e<<endl;
    cout << rootdist(c, getdata(c, a), b) + rootdist(d, getdata(d, a), b) -
                2 * rootdist(e, getdata(e, a), b)
         << "\n";
  }
}


Fixed submission:
#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
#include <vector>
using namespace std;
const int MAXN = 1e5 + 5;
map<int, pair<int, int>> curColours;
vector<map<int, pair<int, int>>> colourData;

vector<pair<int, int>> v1[MAXN];
int level[MAXN];
int parent[MAXN];
int parentedge[MAXN];
int dist[MAXN];
vector<pair<int, int>> edgedata;
int dp[MAXN][24];
bool special[MAXN];
bool checkedi[MAXN];
void dfs1(int curr, int par) {
  for (auto z : v1[curr]) {
    int x = z.first;
    if (x == par) {
      continue;
    }
    level[x] = level[curr] + 1;
    parent[x] = curr;
    parentedge[x] = z.second;
    dist[x] = dist[curr] + edgedata[parentedge[x]].second;
    dfs1(x, curr);
  }
}
void dfs2(int curr, int par) {

  if (special[curr]) {

    colourData[curr] = curColours;
  }
  for (auto x : v1[curr]) {

    if (x.first == par) {
      continue;
    }
    int ind = x.second;
    // cout<<curr<<" "<<x.first<<" "<<ind<<endl;
    int col = edgedata[ind].first;
    int dist = edgedata[ind].second;
    if (!curColours.count(col)) {
      curColours.insert(make_pair(col, make_pair(0, 0)));
    }
    curColours[col].first++;
    curColours[col].second += dist;
    dfs2(x.first, curr);

    curColours[col].first--;
    curColours[col].second -= dist;

    if (curColours[col].first == 0) {
      curColours.erase(col);
    }
  }
}
int lca(int u, int v) {
  if (level[u] < level[v]) {
    swap(u, v);
  }
  for (int i = 20; i >= 0; i--) {
    if (level[u] - (1 << i) >= level[v]) {
      u = dp[u][i];
    }
  }
  if (u == v) {
    return u;
  }
  for (int i = 20; i >= 0; i--) {
    if (dp[u][i] != dp[v][i] && dp[u][i] != -1 && dp[v][i] != -1) {
      u = dp[u][i];
      v = dp[v][i];
    }
  }
  return parent[u];
}
pair<int, int> getdata(int v, int colour) {
  if (v == -1) {
    cout << 1234567 << endl;
    return make_pair(0, 0);
  }
  //   cout<<v<<" "<<parent[v]<<" "<<special[v]<<endl;
  if (special[v]) {
    if (!colourData[v].count(colour)) {
      return make_pair(0, 0);
    } else {
      return colourData[v][colour];
    }
  }

  pair<int, int> prev = getdata(parent[v], colour);
  int ind = parentedge[v];
  if (edgedata[ind].first == colour) {
    prev.first++;
    prev.second += edgedata[ind].second;
  }
  return prev;
}
int rootdist(int curr, pair<int, int> data, int newLen) {
  return dist[curr] - data.second + (data.first * newLen);
}
int main() {
  int n;
  cin >> n;
  int q;
  cin >> q;
  for (int i = 0; i < n - 1; i++) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    edgedata.push_back(make_pair(c, d));
    v1[a].push_back(make_pair(b, i));
    v1[b].push_back(make_pair(a, i));
  }
  parent[1] = -1;
  level[1] = 0;
  dist[1] = 0;
  dfs1(1, -1);
  priority_queue<pair<int, int>> pq;
  for (int i = 1; i <= n; i++) {
    pq.push(make_pair(level[i], i));
  }
  for (int i = 0; i <= n; i++) {
    special[i] = false;
    checkedi[i] = false;
  }

  int cnt = 1;
  special[1] = true;
  checkedi[1] = true;
  while (!pq.empty()) {
    int curr = pq.top().second;
    pq.pop();

    if (checkedi[curr]) {
      continue;
    }
    for (int i = 0; i < 350; i++) {
      curr = parent[curr];
      if (special[curr]) {
        break;
      }
      checkedi[curr] = true;
    }
    special[curr] = true;
  }

  for (int i = 0; i <= n; i++) {
    map<int, pair<int, int>> blank;
    colourData.push_back(blank);
  }
  dfs2(1, -1);
  for (int i = 1; i <= n; i++) {
    for (int j = 0; j <= 20; j++) {
      dp[i][j] = -1;
    }
  }
  for (int i = 1; i <= n; i++) {
    dp[i][0] = parent[i];
  }
  for (int j = 1; j <= 20; j++) {
    for (int i = 1; i <= n; i++) {
      if (dp[i][j - 1] == -1) {
        dp[i][j] = -1;
      } else {
        dp[i][j] = max(dp[i][j], dp[dp[i][j - 1]][j - 1]);
      }
    }
  }
  while (q--) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    int e = lca(c, d);
    // cout<<e<<endl;
    cout << rootdist(c, getdata(c, a), b) + rootdist(d, getdata(d, a), b) -
                2 * rootdist(e, getdata(e, a), b)
         << "\n";
  }
}