Problem p00498 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <algorithm>
#include <functional>
#include <iostream>
#include <queue>
#include <tuple>
#include <vector>
using namespace std;
int N, M, K, Q, a, b, c, dist[101010], alldist[101010], Union[101010];
vector<pair<int, int>> vec[101010], X[101010], X_up[101010], X_down[101010];
vector<tuple<int, int, int>> tup;
vector<int> group[101010], W;
int Before[101010][32], MinDist[101010][32], start[101010], goal[101010];

//-------------------------------Union Find--------------------------------
void unite(int s, int t) {
  int W1 = Union[s], W2 = Union[t];
  if (W1 == W2)
    return;

  if (group[W1].size() < group[W2].size()) {
    for (int i = 0; i < group[W1].size(); i++) {
      group[W2].push_back(group[W1][i]);
      Union[group[W1][i]] = W2;
    }
    group[W1].clear();
    return;
  }
  if (group[W1].size() >= group[W2].size()) {
    for (int i = 0; i < group[W2].size(); i++) {
      group[W1].push_back(group[W2][i]);
      Union[group[W2][i]] = W1;
    }
    group[W2].clear();
    return;
  }
}
bool same(int s, int t) {
  if (Union[s] == Union[t])
    return true;
  return false;
}

//--------------------------------Dijkstra---------------------------------
void dijkstra() {
  priority_queue<pair<int, int>, vector<pair<int, int>>,
                 greater<pair<int, int>>>
      Q;
  for (int i = 0; i < 101010; i++) {
    dist[i] = 999999999;
  }
  for (int i = 0; i < W.size(); i++) {
    dist[W[i]] = 0;
    Q.push(make_pair(0, W[i]));
  }
  while (!Q.empty()) {
    int a1 = Q.top().first, a2 = Q.top().second;
    Q.pop();
    for (int i = 0; i < vec[a2].size(); i++) {
      int to = vec[a2][i].first, cost = vec[a2][i].second;
      if (dist[to] > a1 + cost) {
        dist[to] = a1 + cost;
        Q.push(make_pair(dist[to], to));
      }
    }
  }
}

//---------------------------BreadthFirstSearch----------------------------
void bfs() {
  queue<int> Q;
  for (int i = 0; i < 101010; i++)
    dist[i] = 999999999;
  Q.push(1);
  dist[1] = 0;
  while (!Q.empty()) {
    int a1 = Q.front();
    Q.pop();
    for (int i = 0; i < X[a1].size(); i++) {
      int to = X[a1][i].first;
      if (dist[to] == 999999999) {
        dist[to] = dist[a1] + 1;
        Q.push(to);
      }
    }
  }
}

//----------------------------LCA,Construction-----------------------------
void LCAConstruct() {
  for (int i = 0; i < 101010; i++) {
    for (int j = 0; j < 30; j++)
      Before[i][j] = -1;
  }
  for (int i = 1; i <= N; i++) {
    if (X_up[i].size() >= 1)
      Before[i][0] = X_up[i][0].first;
  }
  for (int i = 1; i < 30; i++) {
    for (int j = 1; j <= N; j++) {
      int Before1 = Before[j][i - 1];
      if (Before1 == -1)
        continue;
      Before[j][i] = Before[Before1][i - 1];
    }
  }
}
void MinDistConstruct() {
  for (int i = 0; i < 101010; i++) {
    for (int j = 0; j < 30; j++)
      MinDist[i][j] = 999999999;
  }
  for (int i = 1; i <= N; i++) {
    if (X_up[i].size() >= 1) {
      MinDist[i][0] = X_up[i][0].second;
    }
  }
  for (int i = 1; i < 30; i++) {
    for (int j = 1; j <= N; j++) {
      int Before1 = Before[j][i - 1];
      if (Before1 == -1)
        continue;
      MinDist[j][i] = min(MinDist[j][i - 1], MinDist[Before1][i - 1]);
    }
  }
}

//--------------------------------Solvings---------------------------------
int Before_solve(int s, int p) {
  int ret2 = s;
  for (int i = 0; i < 30; i++) {
    if (ret2 == -1)
      continue;
    if ((p / (1LL << i)) % 2 == 1)
      ret2 = Before[ret2][i];
  }
  return ret2;
}
int lca(int s, int t) {
  int L = 0, R = 10000000, M;
  if (dist[s] < dist[t])
    t = Before_solve(t, dist[t] - dist[s]);
  if (dist[s] > dist[t])
    s = Before_solve(s, dist[s] - dist[t]);
  if (s == t)
    return s;
  while (true) {
    M = (L + R) / 2;
    int A1 = Before_solve(s, M), A2 = Before_solve(s, M - 1);
    int B1 = Before_solve(t, M), B2 = Before_solve(t, M - 1);
    if (A1 == B1 && A2 != B2)
      return A1;
    if (A1 != B1)
      L = M;
    if (A2 == B2)
      R = M;
  }
}
int rmq(int s, int p) {
  if (p == 0)
    return alldist[s];
  int ret4 = 999999999, pos = s;
  for (int i = 0; i < 30; i++) {
    if ((p / (1 << i)) % 2 == 1) {
      ret4 = min(ret4, MinDist[pos][i]);
      pos = Before[pos][i];
    }
  }
  return ret4;
}
int solve(int s, int t) {
  int lca2 = lca(s, t);
  int dist1 = dist[s] - dist[lca2];
  int dist2 = dist[t] - dist[lca2];
  int rmq1 = rmq(s, dist1);
  int rmq2 = rmq(t, dist2);
  return min(rmq1, rmq2);
}

//----------------------------------Main-----------------------------------
int main() {
  for (int i = 0; i < 101010; i++) {
    group[i].push_back(i);
    Union[i] = i;
  }
  cin >> N >> M >> K >> Q;
  for (int i = 0; i < M; i++) {
    cin >> a >> b >> c;
    vec[a].push_back(make_pair(b, c));
    vec[b].push_back(make_pair(a, c));
  }
  for (int i = 0; i < K; i++) {
    cin >> a;
    W.push_back(a);
  }
  dijkstra();
  for (int i = 1; i <= N; i++)
    alldist[i] = dist[i];
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < vec[i].size(); j++) {
      tup.push_back(
          make_tuple(min(dist[i], dist[vec[i][j].first]), i, vec[i][j].first));
    }
  }
  sort(tup.begin(), tup.end(), greater<tuple<int, int, int>>());
  for (int i = 0; i < tup.size(); i++) {
    int a1 = get<1>(tup[i]), a2 = get<2>(tup[i]);
    if (same(a1, a2) == false) {
      unite(a1, a2);
      X[a1].push_back(make_pair(a2, get<0>(tup[i])));
      X[a2].push_back(make_pair(a1, get<0>(tup[i])));
    }
  }
  bfs();
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < X[i].size(); j++) {
      if (dist[i] < dist[X[i][j].first]) {
        X_down[i].push_back(make_pair(X[i][j].first, X[i][j].second));
      }
      if (dist[i] > dist[X[i][j].first]) {
        X_up[i].push_back(make_pair(X[i][j].first, X[i][j].second));
      }
    }
  }
  LCAConstruct();
  MinDistConstruct();
  for (int i = 0; i < Q; i++) {
    cin >> start[i] >> goal[i];
    cout << solve(start[i], goal[i]) << endl;
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <functional>
#include <iostream>
#include <queue>
#include <tuple>
#include <vector>
using namespace std;
int N, M, K, Q, a, b, c, dist[108108], alldist[108108], Union[108108];
vector<pair<int, int>> vec[108108], X[108108], X_up[108108], X_down[108108];
vector<tuple<int, int, int>> tup;
vector<int> group[108108], W;
int Before[108108][32], MinDist[108108][32], start[108108], goal[108108];

//-------------------------------Union Find--------------------------------
void unite(int s, int t) {
  int W1 = Union[s], W2 = Union[t];
  if (W1 == W2)
    return;

  if (group[W1].size() < group[W2].size()) {
    for (int i = 0; i < group[W1].size(); i++) {
      group[W2].push_back(group[W1][i]);
      Union[group[W1][i]] = W2;
    }
    group[W1].clear();
    return;
  }
  if (group[W1].size() >= group[W2].size()) {
    for (int i = 0; i < group[W2].size(); i++) {
      group[W1].push_back(group[W2][i]);
      Union[group[W2][i]] = W1;
    }
    group[W2].clear();
    return;
  }
}
bool same(int s, int t) {
  if (Union[s] == Union[t])
    return true;
  return false;
}

//--------------------------------Dijkstra---------------------------------
void dijkstra() {
  priority_queue<pair<int, int>, vector<pair<int, int>>,
                 greater<pair<int, int>>>
      Q;
  for (int i = 0; i < 108108; i++) {
    dist[i] = 999999999;
  }
  for (int i = 0; i < W.size(); i++) {
    dist[W[i]] = 0;
    Q.push(make_pair(0, W[i]));
  }
  while (!Q.empty()) {
    int a1 = Q.top().first, a2 = Q.top().second;
    Q.pop();
    for (int i = 0; i < vec[a2].size(); i++) {
      int to = vec[a2][i].first, cost = vec[a2][i].second;
      if (dist[to] > a1 + cost) {
        dist[to] = a1 + cost;
        Q.push(make_pair(dist[to], to));
      }
    }
  }
}

//---------------------------BreadthFirstSearch----------------------------
void bfs() {
  queue<int> Q;
  for (int i = 0; i < 108108; i++)
    dist[i] = 999999999;
  Q.push(1);
  dist[1] = 0;
  while (!Q.empty()) {
    int a1 = Q.front();
    Q.pop();
    for (int i = 0; i < X[a1].size(); i++) {
      int to = X[a1][i].first;
      if (dist[to] == 999999999) {
        dist[to] = dist[a1] + 1;
        Q.push(to);
      }
    }
  }
}

//----------------------------LCA,Construction-----------------------------
void LCAConstruct() {
  for (int i = 0; i < 108108; i++) {
    for (int j = 0; j < 30; j++)
      Before[i][j] = -1;
  }
  for (int i = 1; i <= N; i++) {
    if (X_up[i].size() >= 1)
      Before[i][0] = X_up[i][0].first;
  }
  for (int i = 1; i < 30; i++) {
    for (int j = 1; j <= N; j++) {
      int Before1 = Before[j][i - 1];
      if (Before1 == -1)
        continue;
      Before[j][i] = Before[Before1][i - 1];
    }
  }
}
void MinDistConstruct() {
  for (int i = 0; i < 108108; i++) {
    for (int j = 0; j < 30; j++)
      MinDist[i][j] = 999999999;
  }
  for (int i = 1; i <= N; i++) {
    if (X_up[i].size() >= 1) {
      MinDist[i][0] = X_up[i][0].second;
    }
  }
  for (int i = 1; i < 30; i++) {
    for (int j = 1; j <= N; j++) {
      int Before1 = Before[j][i - 1];
      if (Before1 == -1)
        continue;
      MinDist[j][i] = min(MinDist[j][i - 1], MinDist[Before1][i - 1]);
    }
  }
}

//--------------------------------Solvings---------------------------------
int Before_solve(int s, int p) {
  int ret2 = s;
  for (int i = 0; i < 30; i++) {
    if (ret2 == -1)
      continue;
    if ((p / (1LL << i)) % 2 == 1)
      ret2 = Before[ret2][i];
  }
  return ret2;
}
int lca(int s, int t) {
  int L = 0, R = 10000000, M;
  if (dist[s] < dist[t])
    t = Before_solve(t, dist[t] - dist[s]);
  if (dist[s] > dist[t])
    s = Before_solve(s, dist[s] - dist[t]);
  if (s == t)
    return s;
  while (true) {
    M = (L + R) / 2;
    int A1 = Before_solve(s, M), A2 = Before_solve(s, M - 1);
    int B1 = Before_solve(t, M), B2 = Before_solve(t, M - 1);
    if (A1 == B1 && A2 != B2)
      return A1;
    if (A1 != B1)
      L = M;
    if (A2 == B2)
      R = M;
  }
}
int rmq(int s, int p) {
  if (p == 0)
    return alldist[s];
  int ret4 = 999999999, pos = s;
  for (int i = 0; i < 30; i++) {
    if ((p / (1 << i)) % 2 == 1) {
      ret4 = min(ret4, MinDist[pos][i]);
      pos = Before[pos][i];
    }
  }
  return ret4;
}
int solve(int s, int t) {
  int lca2 = lca(s, t);
  int dist1 = dist[s] - dist[lca2];
  int dist2 = dist[t] - dist[lca2];
  int rmq1 = rmq(s, dist1);
  int rmq2 = rmq(t, dist2);
  return min(rmq1, rmq2);
}

//----------------------------------Main-----------------------------------
int main() {
  for (int i = 0; i < 108108; i++) {
    group[i].push_back(i);
    Union[i] = i;
  }
  cin >> N >> M >> K >> Q;
  for (int i = 0; i < M; i++) {
    cin >> a >> b >> c;
    vec[a].push_back(make_pair(b, c));
    vec[b].push_back(make_pair(a, c));
  }
  for (int i = 0; i < K; i++) {
    cin >> a;
    W.push_back(a);
  }
  dijkstra();
  for (int i = 1; i <= N; i++)
    alldist[i] = dist[i];
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < vec[i].size(); j++) {
      tup.push_back(
          make_tuple(min(dist[i], dist[vec[i][j].first]), i, vec[i][j].first));
    }
  }
  sort(tup.begin(), tup.end(), greater<tuple<int, int, int>>());
  for (int i = 0; i < tup.size(); i++) {
    int a1 = get<1>(tup[i]), a2 = get<2>(tup[i]);
    if (same(a1, a2) == false) {
      unite(a1, a2);
      X[a1].push_back(make_pair(a2, get<0>(tup[i])));
      X[a2].push_back(make_pair(a1, get<0>(tup[i])));
    }
  }
  bfs();
  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < X[i].size(); j++) {
      if (dist[i] < dist[X[i][j].first]) {
        X_down[i].push_back(make_pair(X[i][j].first, X[i][j].second));
      }
      if (dist[i] > dist[X[i][j].first]) {
        X_up[i].push_back(make_pair(X[i][j].first, X[i][j].second));
      }
    }
  }
  LCAConstruct();
  MinDistConstruct();
  for (int i = 0; i < Q; i++) {
    cin >> start[i] >> goal[i];
    cout << solve(start[i], goal[i]) << endl;
  }
  return 0;
}