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

#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

#define INF 1145141919
typedef pair<int, int> P;
typedef pair<int, P> E;

int N, M, K, Q;
int U[18][200000], C[200000], D[200000], T[200000];
int A[200000];
vector<int> L[200000];
vector<P> G[200000];

int find(int x) {
  if (A[x] == x)
    return x;
  return A[x] = find(A[x]);
}

int lca(int a, int b) {
  if (D[a] > D[b])
    swap(a, b);
  for (int k = 0; k < 17; k++) {
    if ((D[a] - D[b]) >> k & 1) {
      b = U[k][b];
    }
  }
  if (a == b)
    return a;
  for (int k = 17; k >= 0; k--) {
    if (U[k][a] != U[k][b]) {
      a = U[k][a];
      b = U[k][b];
    }
  }
  return U[0][a];
}

void unite(int x, int y, int r) {
  x = find(x), y = find(y);
  if (x == y)
    return;
  int p = N++;
  L[p].push_back(x), L[p].push_back(y);
  U[0][x] = p, U[0][y] = p;
  A[x] = p, A[y] = p;
  T[p] = r;
}

int main() {
  cin >> N >> M >> K >> Q;
  for (int i = 0; i < M; i++) {
    int a, b, l;
    cin >> a >> b >> l;
    G[a - 1].push_back(P(l, b - 1));
    G[b - 1].push_back(P(l, a - 1));
  }
  for (int i = 0; i < 2 * N; i++) {
    C[i] = INF;
    A[i] = i;
    T[i] = -1;
    U[0][i] = -1;
  }

  priority_queue<P, vector<P>, greater<P>> q;
  for (int i = 0; i < K; i++) {
    int x;
    cin >> x;
    q.push(P(0, x - 1));
    C[x - 1] = 0;
  }
  while (!q.empty()) {
    int s = q.top().second, r = q.top().first;
    q.pop();
    for (auto p : G[s]) {
      int t = p.second, l = p.first;
      if (C[t] > C[s] + l) {
        C[t] = C[s] + l;
        q.push(P(C[t], t));
      }
    }
  }

  vector<E> g;
  for (int i = 0; i < N; i++) {
    for (auto p : G[i]) {
      int t = p.second;
      g.push_back(E(min(C[i], C[t]), P(i, t)));
    }
  }
  sort(g.begin(), g.end(), greater<E>());

  for (auto e : g) {
    int r = e.first, s = e.second.first, t = e.second.second;
    unite(s, t, r);
  }
  int root = find(0);

  queue<P> bfs;
  bfs.push(P(root, 1));
  D[root] = 1;
  while (!bfs.empty()) {
    int s = bfs.front().first, r = bfs.front().second;
    bfs.pop();
    for (int t : L[s]) {
      D[t] = r + 1;
      bfs.push(P(t, r + 1));
    }
  }
  for (int k = 0; k < 17; k++) {
    for (int v = 0; v < N; v++) {
      if (U[k][v] == -1)
        U[k + 1][v] = -1;
      else
        U[k + 1][v] = U[k][U[k][v]];
    }
  }

  for (int i = 0; i < Q; i++) {
    int s, t;
    cin >> s >> t;
    s--, t--;
    cout << T[lca(s, t)] << "\\n";
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

#define INF 1145141919
typedef pair<int, int> P;
typedef pair<int, P> E;

int N, M, K, Q;
int U[18][200000], C[200000], D[200000], T[200000];
vector<int> L[200000];
vector<P> G[200000];

int find(int x) {
  if (U[0][x] == -1)
    return x;
  return find(U[0][x]);
}

int lca(int a, int b) {
  if (D[a] > D[b])
    swap(a, b);
  for (int k = 0; k < 17; k++) {
    if ((D[a] - D[b]) >> k & 1) {
      b = U[k][b];
    }
  }
  if (a == b)
    return a;
  for (int k = 17; k >= 0; k--) {
    if (U[k][a] != U[k][b]) {
      a = U[k][a];
      b = U[k][b];
    }
  }
  return U[0][a];
}

void unite(int x, int y, int r) {
  x = find(x);
  y = find(y);
  if (x == y)
    return;
  int p = N++;
  L[p].push_back(x), L[p].push_back(y);
  U[0][x] = p, U[0][y] = p;
  T[p] = r;
}

int main() {
  cin >> N >> M >> K >> Q;
  for (int i = 0; i < M; i++) {
    int a, b, l;
    cin >> a >> b >> l;
    G[a - 1].push_back(P(l, b - 1));
    G[b - 1].push_back(P(l, a - 1));
  }
  for (int i = 0; i < 2 * N; i++) {
    C[i] = INF;
    T[i] = -1;
    U[0][i] = -1;
  }

  priority_queue<P, vector<P>, greater<P>> q;
  for (int i = 0; i < K; i++) {
    int x;
    cin >> x;
    q.push(P(0, x - 1));
    C[x - 1] = 0;
  }
  while (!q.empty()) {
    int s = q.top().second, r = q.top().first;
    q.pop();
    for (auto p : G[s]) {
      int t = p.second, l = p.first;
      if (C[t] > C[s] + l) {
        C[t] = C[s] + l;
        q.push(P(C[t], t));
      }
    }
  }

  vector<E> g;
  for (int i = 0; i < N; i++) {
    for (auto p : G[i]) {
      int t = p.second;
      g.push_back(E(min(C[i], C[t]), P(i, t)));
    }
  }
  sort(g.begin(), g.end(), greater<E>());

  for (auto e : g) {
    int r = e.first, s = e.second.first, t = e.second.second;
    unite(s, t, r);
  }
  int root = find(0);

  queue<P> bfs;
  bfs.push(P(root, 1));
  D[root] = 1;
  while (!bfs.empty()) {
    int s = bfs.front().first, r = bfs.front().second;
    bfs.pop();
    for (int t : L[s]) {
      D[t] = r + 1;
      bfs.push(P(t, r + 1));
    }
  }
  for (int k = 0; k < 17; k++) {
    for (int v = 0; v < N; v++) {
      if (U[k][v] == -1)
        U[k + 1][v] = -1;
      else
        U[k + 1][v] = U[k][U[k][v]];
    }
  }

  for (int i = 0; i < Q; i++) {
    int s, t;
    cin >> s >> t;
    s--, t--;
    cout << T[lca(s, t)] << "\\n";
  }
  return 0;
}