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

#include <algorithm>
#include <functional>
#include <iostream>
#include <queue>
#include <tuple>
#include <vector>
using namespace std;
#pragma warning(disable : 4996)

class UnionFind {
public:
  vector<int> par;
  void init(int sz) { par.resize(sz + 1, -1); }
  int root(int pos) {
    if (par[pos] == -1)
      return pos;
    par[pos] = root(par[pos]);
    return par[pos];
  }
  bool same(int u, int v) {
    if (root(u) == root(v))
      return true;
    return false;
  }
  void unite(int u, int v) {
    u = root(u);
    v = root(v);
    if (u != v)
      par[u] = v;
  }
};

long long N, M, K, Q, A[1 << 19], B[1 << 19], C[1 << 19], dist[1 << 18],
    dst[1 << 18], dp[1 << 18][22], dpm[1 << 18][22];
vector<pair<long long, long long>> X[1 << 19], Y[1 << 19];
priority_queue<pair<long long, long long>, vector<pair<long long, long long>>,
               greater<pair<long long, long long>>>
    que;
vector<tuple<long long, int, int>> tup;
bool used[1 << 18];

void dfs(int pos) {
  used[pos] = true;
  for (int i = 0; i < Y[pos].size(); i++) {
    if (used[Y[pos][i].first] == true)
      continue;
    dp[Y[pos][i].first][0] = pos;
    dpm[Y[pos][i].first][0] = Y[pos][i].second;
    dst[Y[pos][i].first] = dst[pos] + 1;
    dfs(Y[pos][i].first);
  }
}

int prevs(int u, int x) {
  for (int i = 21; i >= 0; i--) {
    if (x >= (1 << i)) {
      u = dp[u][i];
      x -= (1 << i);
    }
  }
  return u;
}

long long prevmin(int u, int x) {
  long long ret = (1LL << 60);
  for (int i = 21; i >= 0; i--) {
    if (x >= (1 << i)) {
      ret = min(ret, dpm[u][i]);
      u = dp[u][i];
      x -= (1 << i);
    }
  }
  return ret;
}

int query(int u, int v) {
  int u1 = u, v1 = v;
  if (dst[u] > dst[v])
    swap(u, v);
  v = prevs(v, dst[v] - dst[u]);

  if (u != v) {
    for (int i = 21; i >= 0; i--) {
      if (dp[u][i] != dp[v][i]) {
        u = dp[u][i];
        v = dp[v][i];
      }
    }
    u = dp[u][0];
    v = dp[v][0];
  }

  long long A1 = prevmin(u1, dst[u1] - dst[u]);
  long long A2 = prevmin(v1, dst[v1] - dst[v]);
  return min(A1, A2);
}

int main() {
  cin >> N >> M >> K >> Q;
  for (int i = 0; i < M; i++) {
    scanf("%lld%lld%lld", &A[i], &B[i], &C[i]);
    X[A[i]].push_back(make_pair(B[i], C[i]));
    X[B[i]].push_back(make_pair(A[i], C[i]));
  }
  for (int i = 1; i <= N; i++)
    dist[i] = (1LL << 60);
  for (int i = 0; i < K; i++) {
    long long a;
    scanf("%lld", &a);
    que.push(make_pair(0, a));
    dist[a] = 0;
  }
  while (!que.empty()) {
    int pos = que.top().second;
    que.pop();
    for (int i = 0; i < X[pos].size(); i++) {
      long long to = X[pos][i].first, cost = X[pos][i].second;
      if (dist[to] > dist[pos] + cost) {
        dist[to] = dist[pos] + cost;
        que.push(make_pair(dist[to], to));
      }
    }
  }
  for (int i = 0; i < M; i++)
    tup.push_back(make_tuple(min(dist[A[i]], dist[B[i]]), A[i], B[i]));
  sort(tup.begin(), tup.end());
  reverse(tup.begin(), tup.end());

  UnionFind UF;
  UF.init(N + 2);
  for (int i = 0; i < tup.size(); i++) {
    int p1 = get<1>(tup[i]), p2 = get<2>(tup[i]);
    if (UF.same(p1, p2) == false) {
      UF.unite(p1, p2);
      Y[p1].push_back(make_pair(p2, get<0>(tup[i])));
      Y[p2].push_back(make_pair(p1, get<0>(tup[i])));
    }
  }
  for (int i = 0; i < 22; i++) {
    for (int j = 0; j <= N; j++) {
      dp[i][j] = 0;
      dpm[i][j] = (1LL << 60);
    }
  }
  dfs(1);
  for (int i = 0; i < 21; i++) {
    for (int j = 1; j <= N; j++)
      dp[j][i + 1] = dp[dp[j][i]][i];
  }
  for (int i = 0; i < 21; i++) {
    for (int j = 1; j <= N; j++)
      dpm[j][i + 1] = min(dpm[j][i], dpm[dp[j][i]][i]);
  }

  for (int i = 1; i <= Q; i++) {
    long long x, y;
    scanf("%lld%lld", &x, &y);
    long long Z = query(x, y);
    printf("%lld\\n", Z);
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
int N, M, K, R, a, b, c, d[200000], dist[200000];
vector<pair<int, int>> x[120000];
priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>>
    Q;
int main() {
  cin >> N >> M >> K >> R;
  for (int i = 0; i < M; i++) {
    cin >> a >> b >> c;
    x[a].push_back(make_pair(b, c));
    x[b].push_back(make_pair(a, c));
  }
  for (int i = 0; i < K; i++) {
    cin >> a;
    d[a] = 1;
  }
  for (int i = 1; i <= N; i++) {
    dist[i] = 999999999;
    if (d[i] == 1) {
      dist[i] = 0;
      Q.push(make_pair(0, i));
    }
  }
  while (!Q.empty()) {
    int a1 = Q.top().first, a2 = Q.top().second;
    Q.pop();
    for (int i = 0; i < (int)x[a2].size(); i++) {
      if (dist[x[a2][i].first] > a1 + x[a2][i].second) {
        dist[x[a2][i].first] = a1 + x[a2][i].second;
        Q.push(make_pair(dist[x[a2][i].first], x[a2][i].first));
      }
    }
  }
  for (int i = 1; i <= N; i++)
    d[i] = dist[i];
  for (int i = 0; i < R; i++) {
    cin >> a >> b;
    for (int j = 1; j <= N; j++)
      dist[j] = 0;
    dist[a] = d[a];
    Q.push(make_pair(-1 * dist[a], a));
    while (!Q.empty()) {
      int a1 = -1 * (Q.top().first);
      int a2 = Q.top().second;
      Q.pop();
      for (int j = 0; j < (int)x[a2].size(); j++) {
        int to = x[a2][j].first;
        int cost = min(a1, d[x[a2][j].first]);
        if (dist[to] < cost) {
          dist[to] = cost;
          Q.push(make_pair(-dist[to], to));
        }
      }
    }
    cout << dist[b] << endl;
  }
  return 0;
}