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

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <queue>
#include <vector>
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reg(i, a, b) for (int i = (a); i <= (b); i++)
using namespace std;
#define fir first
#define sec second
typedef long long int lli;
typedef pair<lli, lli> mp;
typedef pair<lli, mp> mmp;

vector<mp> vs[100005];
lli dis[100005] = {};

lli uni[100005];
int find(int p) {
  if (uni[p] == p)
    return p;
  return (uni[p] = find(uni[p]));
}

mp dob[100005][20];
lli dep[100005] = {};

void dfs(int no, int pa, lli co, int dd) {
  dob[no][0] = mp(pa, co);
  dep[no] = dd;
  rep(i, vs[no].size()) {
    lli to = vs[no][i].fir, nco = vs[no][i].sec;
    if (to == pa)
      continue;
    dfs(to, no, nco, dd + 1);
  }
}

void mdob(int p) {
  reg(i, 1, 19) {
    reg(j, 1, p) {
      mp pa = dob[j][i - 1];
      dob[j][i] =
          mp(dob[pa.fir][i - 1].fir, min(dob[pa.fir][i - 1].sec, pa.sec));
    }
  }
}

void succ(int &p, int i, lli &res) {
  res = min(res, dob[p][i].sec);
  p = dob[p][i].fir;
}

lli lca(int p, int q) {
  if (dep[p] < dep[q])
    swap(p, q);
  // printf("st..%d %d\\n",p,q);
  lli res = 10000000000;
  rep(i, 20) {
    if ((dep[p] - dep[q]) & (1 << i)) {
      succ(p, i, res);
    }
  }
  if (p == q)
    return res;

  // printf("tyu..%d %d %lld\\n",p,q,res);
  for (int i = 19; i >= 0; i--) {
    if (dob[p][i].fir != dob[q][i].fir) {
      succ(p, i, res);
      succ(q, i, res);
    }
  }

  if (p != q)
    res = min(res, min(dob[p][0].sec, dob[q][0].sec));
  return res;
}

lli ans[100005];

int m, n, k, qq;
int main(void) {
  scanf("%d%d%d%d", &n, &m, &k, &qq);
  rep(i, m) {
    lli a, b, c;
    scanf("%lld%lld%lld", &a, &b, &c);
    vs[a].push_back(mp(b, c));
    vs[b].push_back(mp(a, c));
  }

  memset(dis, -1, sizeof(dis));
  priority_queue<mp, vector<mp>, greater<mp>> que;
  rep(i, k) {
    lli a;
    scanf("%lld", &a);
    que.push(mp(0, a));
  }

  // dist
  while (!que.empty()) {
    lli no = que.top().sec, nco = que.top().fir;
    que.pop();
    // printf("%lld %lld\\n",no,nco);
    if (dis[no] != -1)
      continue;
    dis[no] = nco;
    rep(i, vs[no].size()) {
      lli to = vs[no][i].fir, co = vs[no][i].sec;
      if (dis[to] != -1)
        continue;
      que.push(mp(nco + co, to));
    }
  }

  /*
  reg(i,1,n){
    printf("d..%lld\\n",dis[i]);
  }*/
  // kra
  reg(i, 1, n) uni[i] = i;
  vector<mmp> kra;
  reg(i, 1, n) {
    rep(j, vs[i].size()) {
      lli to = vs[i][j].fir;
      if (to >= i)
        continue;
      kra.push_back(mmp(min(dis[i], dis[to]), mp(i, to)));
    }
    vs[i].clear();
  }
  sort(kra.begin(), kra.end(), greater<mmp>());

  rep(i, kra.size()) {
    mmp pa = kra[i];
    lli a = pa.sec.fir, b = pa.sec.sec, c = pa.fir;
    if (find(a) == find(b))
      continue;
    uni[uni[a]] = uni[b];
    // printf("kra..%lld %lld %lld\\n",a,b,c);
    vs[a].push_back(mp(b, c));
    vs[b].push_back(mp(a, c));
  }

  // dob
  dfs(1, 0, -1, 1);

  /*
  reg(i,1,n){
    printf("dep..%lld ed..%d %lld %lld\\n",dep[i],i,dob[i][0].fir,dob[i][0].sec);
  }*/
  mdob(n);

  rep(i, qq) {
    int a, b;
    scanf("%d%d", &a, &b);
    printf("%lld\\n", lca(a, b));
  }
  // rep(i,qq)printf("%lld\\n",ans[i]);
  return 0;
}

Error: TLE

Fixed submission:

#include <cstdio>
#include <cstring>
#include <queue>
#include <vector>
using namespace std;

#define MAX 100005
#define LMAX 20
#define DMAX 200000005
typedef pair<int, int> mp;
typedef pair<int, mp> mmp;
#define fir first
#define sec second
#define rep(i, s, n) for (int i = (s); i < (n); i++)
#define min(a, b) (((a) > (b)) ? (b) : (a))
#define max(a, b) (((a) < (b)) ? (b) : (a))

int n, m, k, q;

vector<mp> ver[MAX];
int dist[MAX];
// dist
vector<mp> tree[MAX];
int uni[MAX];
void uini() {
  rep(i, 0, n + 1) { uni[i] = i; }
}
int upar(int p) {
  if (uni[p] == p)
    return p;
  else
    return (uni[p] = upar(uni[p]));
}
bool usame(int p, int q) {
  if (upar(p) == upar(q))
    return true;
  else
    return false;
}
void umer(int p, int q) { uni[uni[p]] = uni[q]; }
// unifo

mp dob[MAX][LMAX];
int dep[MAX] = {};
void dfs(int no, int pa, int di, int de) {
  dep[no] = de;
  dob[no][0].sec = pa;
  dob[no][0].fir = di;
  rep(i, 0, tree[no].size()) {
    if (dep[tree[no][i].sec] == 0) {
      dfs(tree[no][i].sec, no, tree[no][i].fir, de + 1);
    }
  }
}
// dob

int por(int fr, int by, int &ans) {
  ans = min(ans, dob[fr][by].fir);
  return dob[fr][by].sec;
}
int lca(int p, int q) {
  int ans = MAX;
  if (dep[p] < dep[q]) {
    int a = p;
    p = q;
    q = a;
  }
  if (dep[p] != dep[q]) {
    int dif = dep[p] - dep[q], a = 0;
    while (dif) {
      if ((dif & 1) == 1) {
        p = por(p, a, ans);
      }
      dif /= 2;
      a++;
    }
  }
  if (p == 1 && q == 1)
    return ans;
  while (p != q) {
    rep(i, 0, LMAX) {
      if (dob[p][i].sec == dob[q][i].sec) {
        if (i == 0) {
          p = por(p, 0, ans);
          q = por(q, 0, ans);
          break;
        }
        p = por(p, i - 1, ans);
        q = por(q, i - 1, ans);
        break;
      }
    }
  }
  return ans;
}
// lca&res

int main(void) {
  memset(dist, -1, sizeof(dist));
  scanf("%d%d%d%d", &n, &m, &k, &q);
  rep(i, 0, m) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);
    ver[a].push_back(mp(c, b));
    ver[b].push_back(mp(c, a));
  }
  priority_queue<mp, vector<mp>, greater<mp>> que;
  rep(i, 0, k) {
    int a;
    scanf("%d", &a);
    que.push(mp(0, a));
  }

  while (!que.empty()) {
    int to = que.top().sec, co = que.top().fir;
    que.pop();
    if (dist[to] != -1)
      continue;
    dist[to] = co;
    rep(i, 0, ver[to].size()) {
      que.push(mp(co + ver[to][i].fir, ver[to][i].sec));
    }
  }
  // dist

  uini();
  priority_queue<mmp> krus;
  rep(i, 1, n + 1) {
    rep(j, 0, ver[i].size()) {
      int to = ver[i][j].sec, co = ver[i][j].fir;
      krus.push(mmp(min(dist[i], dist[to]), mp(i, to)));
    }
  }
  while (!krus.empty()) {
    int a = krus.top().sec.fir, b = krus.top().sec.sec, c = krus.top().fir;
    krus.pop();
    if (!usame(a, b)) {
      umer(a, b);
      tree[a].push_back(mp(c, b));
      tree[b].push_back(mp(c, a));
    }
  }
  // unifo

  dfs(1, -1, DMAX, 1);
  rep(i, 1, LMAX) {
    rep(j, 0, n) {
      mp np = dob[j][i - 1];
      if (np.sec == -1) {
        dob[j][i] = np;
        continue;
      }
      mp pa = dob[np.sec][i - 1];
      dob[j][i].sec = pa.sec;
      dob[j][i].fir = min(np.fir, pa.fir);
    }
  }
  // dob
  rep(i, 0, q) {
    int a, b;
    scanf("%d%d", &a, &b);
    printf("%d\\n", lca(a, b));
  }
  // lca&res

  return 0;
}