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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
typedef double db;

#define fr first
#define sc second
#define pb push_back
#define rep(i, x) for (ll i = 0; i < x; i++)
#define rep1(i, x) for (ll i = 1; i <= x; i++)
#define rrep(i, x) for (ll i = x - 1; i >= 0; i--)
#define rrep1(i, x) for (ll i = x; i > 0; i--)

const ll INF = 0xffffffffff;
struct edge {
  ll to, cost;
};

vector<edge> G1[100000];
vector<ll> G[100000];
ll N, M, K, Q;
ll A[200000], B[200000], L[200000];
priority_queue<P, vector<P>, greater<P>> q;
ll d[100000];
vector<ll> ds;
P qs[100000];
ll par[20][100000];
ll parmin[20][100000];
ll depth[100000];
ll logn;
ll vmin[200000];
ll ss[200000];
ll uf[100000];
void init(ll n) {
  rep(i, n) uf[i] = i;
  return;
}

ll find(ll x) {
  if (uf[x] == x)
    return x;
  return uf[x] = find(uf[x]);
}

bool same(ll x, ll y) { return find(x) == find(y); }

void unite(ll x, ll y) {
  x = find(x);
  y = find(y);
  uf[x] = y;
  return;
}

void crafttree(ll v, ll from) {
  par[0][v] = from;
  depth[v] = depth[from] + 1;
  parmin[0][v] = min(d[v], d[from]);
  rep(i, G[v].size()) {
    if (G[v][i] != from)
      crafttree(G[v][i], v);
  }
  return;
}

int main() {
  cin >> N >> M >> K >> Q;
  for (ll j = 1; j < N; j *= 2) {
    logn++;
  }
  // puts("1");
  rep(i, M) {
    cin >> A[i] >> B[i] >> L[i];
    A[i]--;
    B[i]--;
    G1[A[i]].pb(edge{B[i], L[i]});
    G1[B[i]].pb(edge{A[i], L[i]});
  }
  fill(d, d + N, INF);
  rep(i, K) {
    ll x;
    cin >> x;
    --x;
    d[x] = 0;
    q.push(P(d[x], x));
  }
  while (!q.empty()) {
    P p = q.top();
    q.pop();
    ll x = p.sc;
    if (d[x] < p.fr)
      continue;
    rep(i, G1[x].size()) {
      edge e = G1[x][i];
      if (d[e.to] > d[x] + e.cost) {
        d[e.to] = d[x] + e.cost;
        q.push(P(d[e.to], e.to));
      }
    }
  }
  // puts("2");
  rep(i, Q) {
    cin >> qs[i].fr >> qs[i].sc;
    qs[i].fr--;
    qs[i].sc--;
  }
  rep(i, N) ds.pb(d[i]);
  sort(ds.begin(), ds.end());
  ds.erase(unique(ds.begin(), ds.end()), ds.end());
  rep(i, N) { d[i] = lower_bound(ds.begin(), ds.end(), d[i]) - ds.begin(); }
  // priority_queue<P> q1;//?????????????°???¨id
  init(N);
  // rep(i,M){
  // puts("3");
  //	q1.push(P(min(d[A[i]],d[B[i]]),i));
  // }
  rep(i, M) {
    ss[i] = i;
    vmin[i] = min(d[A[i]], d[B[i]]);
  }
  sort(ss, ss + M, [&](const ll &a, const ll &b) { return vmin[a] > vmin[b]; });
  // while(!q1.empty()){
  //	ll id=q1.top().sc;
  //	q1.pop();
  // puts("4");
  rep(i, M) {
    ll id = ss[i];
    // cout<<ss[i]<<endl;
    if (!same(A[id], B[id])) {
      G[A[id]].pb(B[id]);
      G[B[id]].pb(A[id]);
      unite(A[id], B[id]);
    }
  }
  par[0][0] = -1;
  parmin[0][0] = d[0];
  depth[0] = 0;
  rep(i, G[0].size()) crafttree(G[0][i], 0);
  // puts("5");
  for (ll i = 1; i <= logn; i++) {
    rep(k, N) {
      if (par[i - 1][k] != -1)
        par[i][k] = par[i - 1][par[i - 1][k]];
      else
        par[i][k] = -1;
      if (par[i - 1][k] != -1)
        parmin[i][k] = min(parmin[i - 1][k], parmin[i - 1][par[i - 1][k]]);
      else
        parmin[i][k] = parmin[i - 1][k];
    }
  }
  ll a, b, ans, cnt;
  rep(i, Q) {
    a = qs[i].fr;
    b = qs[i].sc;
    if (depth[a] < depth[b])
      swap(a, b);
    ans = min(d[a], d[b]);
    ll D = depth[a] - depth[b]; // D 100111
    cnt = 0;
    while (D > 0) {
      if (D & 1) {
        ans = min(ans, parmin[cnt][a]);
        a = par[cnt][a];
      }
      D >>= 1;
      cnt++;
    }
    // if(depth[a]!=depth[b]) cout<<"bag"<<endl;
    if (a == b) {
      // cout<<"a="<<a<<" b="<<b<<endl;
      cout << ds[ans] << endl;
      continue;
    }
    rrep(i, logn + 1) {
      if (par[i][a] != par[i][b]) {
        ans = min(ans, min(parmin[i][a], parmin[i][b]));
        a = par[i][a];
        b = par[i][b];
      }
    }
    cout << ds[min(ans, min(parmin[0][a], parmin[0][b]))] << endl;
  }
  // cout<<"a"<<endl;
  // cout<<depth[0]<<endl<<endl;
  /*rep(i,N){
          rep(j,N){
                  if(depth[j]==i)
  cout<<j+1<<"d="<<d[j]<<":par="<<par[0][j]+1<<":m="<<parmin[0][j]<<" ";
          }
          cout<<endl;
  }*/
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
typedef double db;

#define fr first
#define sc second
#define pb push_back
#define rep(i, x) for (ll i = 0; i < x; i++)
#define rep1(i, x) for (ll i = 1; i <= x; i++)
#define rrep(i, x) for (ll i = x - 1; i >= 0; i--)
#define rrep1(i, x) for (ll i = x; i > 0; i--)

const ll INF = 0xffffffffff;
struct edge {
  ll to, cost;
};

vector<edge> G1[100000];
vector<ll> G[100000];
ll N, M, K, Q;
ll A[200000], B[200000], L[200000];
priority_queue<P, vector<P>, greater<P>> q;
ll d[100000];
vector<ll> ds;
P qs[100000];
ll par[20][100000];
ll parmin[20][100000];
ll depth[100000];
ll logn;
ll vmin[200000];
ll ss[200000];
ll uf[100000];
void init(ll n) {
  rep(i, n) uf[i] = i;
  return;
}

ll find(ll x) {
  if (uf[x] == x)
    return x;
  return x = find(uf[x]);
}

bool same(ll x, ll y) { return find(x) == find(y); }

void unite(ll x, ll y) {
  x = find(x);
  y = find(y);
  uf[x] = y;
  return;
}

void crafttree(ll v, ll from) {
  par[0][v] = from;
  depth[v] = depth[from] + 1;
  parmin[0][v] = min(d[v], d[from]);
  rep(i, G[v].size()) {
    if (G[v][i] != from)
      crafttree(G[v][i], v);
  }
  return;
}

int main() {
  cin >> N >> M >> K >> Q;
  for (ll j = 1; j < N; j *= 2) {
    logn++;
  }
  rep(i, M) {
    cin >> A[i] >> B[i] >> L[i];
    A[i]--;
    B[i]--;
    G1[A[i]].pb(edge{B[i], L[i]});
    G1[B[i]].pb(edge{A[i], L[i]});
  }
  fill(d, d + N, INF);
  rep(i, K) {
    ll x;
    cin >> x;
    --x;
    d[x] = 0;
    q.push(P(d[x], x));
  }
  while (!q.empty()) {
    P p = q.top();
    q.pop();
    ll x = p.sc;
    if (d[x] < p.fr)
      continue;
    rep(i, G1[x].size()) {
      edge e = G1[x][i];
      if (d[e.to] > d[x] + e.cost) {
        d[e.to] = d[x] + e.cost;
        q.push(P(d[e.to], e.to));
      }
    }
  }
  rep(i, Q) {
    cin >> qs[i].fr >> qs[i].sc;
    qs[i].fr--;
    qs[i].sc--;
  }
  rep(i, N) ds.pb(d[i]);
  sort(ds.begin(), ds.end());
  ds.erase(unique(ds.begin(), ds.end()), ds.end());
  rep(i, N) { d[i] = lower_bound(ds.begin(), ds.end(), d[i]) - ds.begin(); }
  // priority_queue<P> q1;//?????????????°???¨id
  init(N);
  // rep(i,M){
  //	q1.push(P(min(d[A[i]],d[B[i]]),i));
  // }
  rep(i, M) {
    ss[i] = i;
    vmin[i] = min(d[A[i]], d[B[i]]);
  }
  sort(ss, ss + M, [&](const ll &a, const ll &b) { return vmin[a] > vmin[b]; });
  // while(!q1.empty()){
  //	ll id=q1.top().sc;
  //	q1.pop();
  rep(i, M) {
    ll id = ss[i];
    // cout<<ss[i]<<endl;
    if (!same(A[id], B[id])) {
      G[A[id]].pb(B[id]);
      G[B[id]].pb(A[id]);
      unite(A[id], B[id]);
    }
  }
  par[0][0] = -1;
  parmin[0][0] = d[0];
  depth[0] = 0;
  rep(i, G[0].size()) crafttree(G[0][i], 0);
  for (ll i = 1; i <= logn; i++) {
    rep(k, N) {
      if (par[i - 1][k] != -1)
        par[i][k] = par[i - 1][par[i - 1][k]];
      else
        par[i][k] = -1;
      if (par[i - 1][k] != -1)
        parmin[i][k] = min(parmin[i - 1][k], parmin[i - 1][par[i - 1][k]]);
      else
        parmin[i][k] = parmin[i - 1][k];
    }
  }
  ll a, b, ans, cnt;
  rep(i, Q) {
    a = qs[i].fr;
    b = qs[i].sc;
    if (depth[a] < depth[b])
      swap(a, b);
    ans = min(d[a], d[b]);
    ll D = depth[a] - depth[b]; // D 100111
    cnt = 0;
    while (D > 0) {
      if (D & 1) {
        ans = min(ans, parmin[cnt][a]);
        a = par[cnt][a];
      }
      D >>= 1;
      cnt++;
    }
    // if(depth[a]!=depth[b]) cout<<"bag"<<endl;
    if (a == b) {
      // cout<<"a="<<a<<" b="<<b<<endl;
      cout << ds[ans] << endl;
      continue;
    }
    rrep(i, logn + 1) {
      if (par[i][a] != par[i][b]) {
        ans = min(ans, min(parmin[i][a], parmin[i][b]));
        a = par[i][a];
        b = par[i][b];
      }
    }
    cout << ds[min(ans, min(parmin[0][a], parmin[0][b]))] << endl;
  }
  // cout<<"a"<<endl;
  // cout<<depth[0]<<endl<<endl;
  /*rep(i,N){
          rep(j,N){
                  if(depth[j]==i)
  cout<<j+1<<"d="<<d[j]<<":par="<<par[0][j]+1<<":m="<<parmin[0][j]<<" ";
          }
          cout<<endl;
  }*/
}