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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ll long long
using namespace std;
const int MAXN = 200001;
int n, m, k, q;
vector<int> festival;
vector<pair<int, int>> adj[MAXN];
int dist[MAXN];
vector<int> weight;
int parent[MAXN], low[MAXN], high[MAXN];
vector<int> query[MAXN];
int s[MAXN], e[MAXN];
vector<int> group[MAXN];
void dijkstra() {
  priority_queue<pair<int, int>> pq;
  for (int i = 1; i < MAXN; i++)
    dist[i] = 2e9;
  for (int i = 0; i < festival.size(); i++) {
    int x = festival[i];
    dist[x] = 0;
    pq.push(make_pair(0, x));
  }

  while (!pq.empty()) {
    int here = pq.top().second;
    int cost = -pq.top().first;
    pq.pop();
    if (cost > dist[here])
      continue;
    for (int i = 0; i < adj[here].size(); i++) {
      int there = adj[here][i].first;
      int w = adj[here][i].second;
      if (dist[there] > dist[here] + w) {
        dist[there] = dist[here] + w;
        pq.push(make_pair(-dist[there], there));
      }
    }
  }
  for (int i = 1; i <= n; i++) {
    weight.push_back(dist[i]);
  }
  sort(weight.begin(), weight.end());
  weight.erase(unique(weight.begin(), weight.end()), weight.end());
  for (int i = 0; i < MAXN; i++) {
    low[i] = 0;
    high[i] = weight.size() - 1;
  }
  for (int i = 1; i <= n; i++) {
    int lb =
        lower_bound(weight.begin(), weight.end(), dist[i]) - weight.begin();
    group[lb].push_back(i);
  }
}
void init() {
  for (int i = 1; i <= n; i++)
    parent[i] = i;
  for (int i = 0; i < MAXN; i++) {
    query[i].clear();
  }
}
int find(int u) { return u == parent[u] ? u : parent[u] = find(parent[u]); }
void mrg(int u, int v) {
  //    assert(find(u) > 0 && find(v) > 0);
  u = find(u);
  v = find(v);
  //    assert(u > 0 && v > 0);
  if (u == v)
    return;
  parent[u] = v;
}
int iter = 0;
bool pbs() {
  iter++;
  if (iter == 30) {
    int abc = 3;
  }
  init();
  int update = 0;
  for (int i = 1; i <= q; i++) {
    if (low[i] <= high[i]) {
      if (i == 33482) {
        int abc = 3;
      }
      int mid = (low[i] + high[i]) / 2;
      query[mid].push_back(i);
      update = 1;
    }
  }
  for (int i = weight.size() - 1; i >= 0; i--) {
    // weight가 i인 녀석들을 다 집어넣고 union
    for (int id : group[i]) {
      for (pair<int, int> it : adj[id]) {
        int nxt = it.first;
        if (dist[nxt] >= weight[i] && dist[id] >= weight[i]) {
          mrg(nxt, id);
        }
      }
    }
    for (int idx : query[i]) {
      if (idx == 33482) {
        int abc = 4;
      }
      int ss = s[idx];
      int ee = e[idx];
      assert(ss > 0 && ee > 0);
      if (idx == 33482) {
        int abc = 3;
      }
      if (find(ss) == find(ee)) {
        low[idx] = i + 1;
      } else
        high[idx] = i - 1;
    }
  }
  return update;
}
int main() {
  // freopen("/Users/gangseong-yong/Downloads/2012-ho-data/2012-ho-t5/in/09-02.txt","r",stdin);
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> n >> m >> k >> q;
  for (int i = 1; i <= m; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    adj[u].push_back({v, w});
    adj[v].push_back({u, w});
  }
  for (int i = 1; i <= k; i++) {
    int f;
    cin >> f;
    festival.push_back(f);
  }
  dijkstra();
  for (int i = 1; i <= q; i++) {
    cin >> s[i] >> e[i];
  }
  while (pbs()) {
  }
  for (int i = 1; i <= q; i++) {
    cout << weight[low[i] - 1] << "\\n";
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ll long long
using namespace std;
const int MAXN = 200002;
const int MAXB = 200002;
int n, m, k, q;
vector<int> festival;
vector<pair<int, int>> adj[MAXB];
int dist[MAXB];
vector<int> weight;
int parent[MAXB], low[MAXN], high[MAXN];
vector<int> query[MAXN];
int s[MAXN], e[MAXN];
vector<int> group[MAXN];
void dijkstra() {
  priority_queue<pair<int, int>> pq;
  for (int i = 1; i < MAXN; i++)
    dist[i] = 2e9;
  for (int i = 0; i < festival.size(); i++) {
    int x = festival[i];
    dist[x] = 0;
    pq.push(make_pair(0, x));
  }

  while (!pq.empty()) {
    int here = pq.top().second;
    int cost = -pq.top().first;
    pq.pop();
    if (cost > dist[here])
      continue;
    for (int i = 0; i < adj[here].size(); i++) {
      int there = adj[here][i].first;
      int w = adj[here][i].second;
      if (dist[there] > dist[here] + w) {
        dist[there] = dist[here] + w;
        pq.push(make_pair(-dist[there], there));
      }
    }
  }
  for (int i = 1; i <= n; i++) {
    weight.push_back(dist[i]);
  }
  sort(weight.begin(), weight.end());
  weight.erase(unique(weight.begin(), weight.end()), weight.end());
  for (int i = 0; i < MAXN; i++) {
    low[i] = 0;
    high[i] = weight.size() - 1;
  }
  for (int i = 1; i <= n; i++) {
    int lb =
        lower_bound(weight.begin(), weight.end(), dist[i]) - weight.begin();
    group[lb].push_back(i);
  }
}
void init() {
  for (int i = 1; i <= n; i++)
    parent[i] = i;
  for (int i = 1; i < MAXN; i++) {
    query[i].clear();
  }
}
int find(int u) { return u == parent[u] ? u : parent[u] = find(parent[u]); }
void mrg(int u, int v) {
  //    assert(find(u) > 0 && find(v) > 0);
  u = find(u);
  v = find(v);
  //    assert(u > 0 && v > 0);
  if (u == v)
    return;
  parent[u] = v;
}
bool pbs() {
  init();
  int update = 0;
  for (int i = 1; i <= q; i++) {
    if (low[i] <= high[i]) {
      query[(low[i] + high[i]) / 2].push_back(i);
      update = 1;
    }
  }
  for (int i = weight.size() - 1; i >= 0; i--) {
    // weight가 i인 녀석들을 다 집어넣고 union
    for (int id : group[i]) {
      for (pair<int, int> it : adj[id]) {
        int nxt = it.first;
        if (dist[nxt] >= weight[i] && dist[id] >= weight[i]) {
          mrg(nxt, id);
        }
      }
    }
    for (int idx : query[i]) {
      int ss = s[idx];
      int ee = e[idx];
      if (find(ss) == find(ee)) {
        low[idx] = i + 1;
      } else
        high[idx] = i - 1;
    }
  }
  return update;
}
int main() {
  // freopen("/Users/gangseong-yong/Desktop/KangSungYong/KangSungYong/in.in","r",stdin);
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> n >> m >> k >> q;
  for (int i = 1; i <= m; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    adj[u].push_back({v, w});
    adj[v].push_back({u, w});
  }
  for (int i = 1; i <= k; i++) {
    int f;
    cin >> f;
    festival.push_back(f);
  }
  dijkstra();
  for (int i = 1; i <= q; i++) {
    cin >> s[i] >> e[i];
  }
  while (pbs()) {
  }
  for (int i = 1; i <= q; i++) {
    cout << weight[low[i] - 1] << "\\n";
  }
  return 0;
}