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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

struct UnionFind {
  vector<int> par;
  int cnt;
  UnionFind(int size_) : par(size_, -1), cnt(size_) {}
  void unite(int x, int y) {
    if ((x = find(x)) != (y = find(y))) {
      if (par[y] < par[x])
        swap(x, y);
      par[x] += par[y];
      par[y] = x;
      cnt--;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int find(int x) { return par[x] < 0 ? x : par[x] = find(par[x]); }
  int size(int x) { return -par[find(x)]; }
  int size() { return cnt; }
};

const int INF = 1e9;

vector<int> dist(100000);
typedef pair<int, int> P;

struct edge {
  int from, to, d;
  bool operator<(const edge &x) const { return d > x.d; }
};

vector<edge> G[100000];
edge es[200000];
int a[100000];
int b[100000];
set<int> qs[100000];
int ans[100000];

int main() {
  int N, M, K, Q;
  scanf("%d%d%d%d", &N, &M, &K, &Q);

  for (int i = 0; i < M; i++) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);
    a--, b--;
    G[a].push_back({a, b, c});
    G[b].push_back({b, a, c});
    es[i] = {a, b, -1};
  }

  fill(dist.begin(), dist.end(), INF);
  priority_queue<P, vector<P>, greater<P>> q;
  for (int i = 0; i < K; i++) {
    int a;
    scanf("%d", &a);
    a--;
    dist[a] = 0;
    q.push({0, a});
  }
  while (q.size()) {
    int v = q.top().second, d = q.top().first;
    q.pop();
    if (dist[v] < d)
      continue;
    for (auto &e : G[v]) {
      if (dist[e.to] > d + e.d) {
        dist[e.to] = d + e.d;
        q.push({d + e.d, e.to});
      }
    }
  }

  for (int i = 0; i < M; i++) {
    int u = es[i].from, v = es[i].to;
    es[i].d = min(dist[u], dist[v]);
  }
  sort(es, es + M);

  for (int i = 0; i < Q; i++) {
    scanf("%d%d", &a[i], &b[i]);
    a[i]--, b[i]--;
    qs[a[i]].insert(i);
    qs[b[i]].insert(i);
  }

  UnionFind uf(N);
  for (int i = 0; i < M; i++) {
    int u = es[i].from, v = es[i].to;
    if (uf.same(u, v))
      continue;
    int up = uf.find(u), vp = uf.find(v);
    if (uf.size(up) < uf.size(vp))
      swap(up, vp);
    uf.unite(u, v);
    for (auto p : qs[vp]) {
      if (uf.same(a[p], b[p])) {
        ans[p] = es[i].d;
        qs[up].erase(p);
      } else {
        qs[up].insert(p);
      }
    }
  }
  for (int i = 0; i < Q; i++) {
    printf("%d\\n", ans[i]);
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

struct UnionFind {
  vector<int> par;
  int cnt;
  UnionFind(int size_) : par(size_, -1), cnt(size_) {}
  void unite(int x, int y) {
    if ((x = find(x)) != (y = find(y))) {
      if (par[y] < par[x])
        swap(x, y);
      par[x] += par[y];
      par[y] = x;
      cnt--;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int find(int x) { return par[x] < 0 ? x : par[x] = find(par[x]); }
  int size(int x) { return -par[find(x)]; }
  int size() { return cnt; }
};

const int MAX_V = 100010;
const int MAX_LOG_V = 21;
const int INF = 1e9;

int dist[100000];
class LCA {
public:
  vector<int> G[MAX_V];
  int root;

  int parent[MAX_LOG_V][MAX_V];
  int depth[MAX_V];
  int mind[MAX_LOG_V][MAX_V];

  void dfs(int v, int p, int d) {
    parent[0][v] = p;
    depth[v] = d;
    for (int i = 0; i < G[v].size(); i++) {
      if (G[v][i] != p)
        dfs(G[v][i], v, d + 1);
    }
  }

  void add_edge(int a, int b) {
    G[a].push_back(b);
    G[b].push_back(a);
  }

  void init(int V, int root = 0) {
    this->root = root;
    dfs(root, -1, 0);
    for (int i = 0; i < V; i++) {
      for (int j = 0; j < MAX_LOG_V; j++) {
        mind[j][i] = INF;
      }
      mind[0][i] = dist[i];
    }
    for (int k = 0; k + 1 < MAX_LOG_V; k++) {
      for (int v = 0; v < V; v++) {
        if (parent[k][v] < 0)
          parent[k + 1][v] = -1;
        else {
          int p = parent[k][parent[k][v]];
          parent[k + 1][v] = p;
          mind[k + 1][v] = min({mind[k][v], mind[k][parent[k][v]], dist[p]});
        }
      }
    }
  }

  int lca(int u, int v) {
    if (depth[u] > depth[v])
      swap(u, v);
    for (int k = 0; k < MAX_LOG_V; k++) {
      if ((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }
    if (u == v)
      return u;
    for (int k = MAX_LOG_V - 1; k >= 0; k--) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }
};

typedef pair<int, int> P;

struct edge {
  int from, to, d;
  bool operator<(const edge &x) const { return d > x.d; }
};

vector<edge> G[100000];
edge es[200000];

LCA lca;

int main() {
  int N, M, K, Q;
  scanf("%d%d%d%d", &N, &M, &K, &Q);

  for (int i = 0; i < M; i++) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);
    a--, b--;
    G[a].push_back({a, b, c});
    G[b].push_back({b, a, c});
    es[i] = {a, b, -1};
  }

  fill((int *)begin(dist), (int *)end(dist), INF);
  priority_queue<P> q;
  for (int i = 0; i < K; i++) {
    int a;
    scanf("%d", &a);
    a--;
    dist[a] = 0;
    q.push({0, a});
  }
  while (q.size()) {
    int v = q.top().second, d = q.top().first;
    q.pop();
    if (dist[v] < d)
      continue;
    for (auto e : G[v]) {
      if (dist[e.to] > d + e.d) {
        dist[e.to] = d + e.d;
        q.push({d + e.d, e.to});
      }
    }
  }
  cout << -1 << endl;
  return 0;

  for (int i = 0; i < M; i++) {
    int u = es[i].from, v = es[i].to;
    es[i].d = min(dist[u], dist[v]);
  }
  sort(es, es + M);

  UnionFind uf(N);
  for (int i = 0; i < M; i++) {
    int u = es[i].from, v = es[i].to;
    if (uf.same(u, v))
      continue;
    uf.unite(u, v);
    lca.add_edge(u, v);
  }
  lca.init(N);

  while (Q--) {
    int a, b;
    cin >> a >> b;
    a--, b--;

    int x = lca.lca(a, b);
    int s[2] = {a, b};
    int ans = INF;
    for (int i = 0; i < 2; i++) {
      int v = s[i];
      int d = lca.depth[v] - lca.depth[x];
      for (int l = 0; l < 20; l++) {
        if (d >> l & 1) {
          ans = min(ans, lca.mind[l][v]);
          v = lca.parent[l][v];
        }
      }
    }
    printf("%d\\n", ans);
  }
}