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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

typedef long long ll;
using namespace std;

#define mod 1000000007
#define INF 1000000000
#define LLINF 2000000000000000000LL

#define SIZE 100000
#define MAXLOG 18

int n, m, k, q;
vector<pair<int, int>> way[SIZE], tree[SIZE];
vector<pair<int, pair<int, int>>> way2;
bool visit[SIZE];
int dis[SIZE];

int ufdata[SIZE];

int ufroot(int a) {
  if (ufdata[a] == -1)
    return a;
  return ufdata[a] = ufroot(ufdata[a]);
}

bool ufset(int a, int b) {
  a = ufroot(a);
  b = ufroot(b);

  if (a == b)
    return false;
  ufdata[a] = b;
  return true;
}

// LCA
int parent[MAXLOG][SIZE];
int min_dis[MAXLOG][SIZE];
int depth[SIZE];

void lcadfs(int now, int h = 0, int back = -1) {
  parent[0][now] = back;
  depth[now] = h;

  for (int i = 0; i < tree[now].size(); i++) {
    if (tree[now][i].first != back) {
      min_dis[0][tree[now][i].first] = tree[now][i].second;
      lcadfs(tree[now][i].first, h + 1, now);
    }
  }
}

int main() {
  int a, b, c;

  scanf("%d%d%d%d", &n, &m, &k, &q);

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

    way2.push_back({0, {a, b}});
  }

  // dijkstra
  priority_queue<pair<int, int>> pq;

  for (int i = 0; i < k; i++) {
    scanf("%d", &a);
    a--;
    pq.push({0, a});
  }

  while (pq.size()) {
    pair<int, int> p = pq.top();
    pq.pop();

    int now = p.second;
    int now_dis = p.first;

    if (visit[now])
      continue;
    visit[now] = true;

    dis[now] = -now_dis;

    for (int i = 0; i < way[now].size(); i++) {
      pq.push({now_dis - way[now][i].second, way[now][i].first});
    }
  }

  //?????§??¨?????¨

  for (int i = 0; i < way2.size(); i++) {
    way2[i].first =
        -1 * min(dis[way2[i].second.first], dis[way2[i].second.second]);
  }

  for (int i = 0; i < n; i++)
    ufdata[i] = -1;

  sort(way2.begin(), way2.end());

  assert(m == (int)way2.size());

  for (int i = 0; i < m; i++) {
    a = way2[i].second.first;
    b = way2[i].second.second;
    c = -way2[i].first;

    if (ufset(a, b)) {
      tree[a].push_back({b, c});
      tree[b].push_back({a, c});
    }
  }

  // LCA
  for (int i = 0; i < MAXLOG; i++) {
    for (int j = 0; j < n; j++) {
      parent[i][j] = -1;
      min_dis[i][j] = INF;
    }
  }

  lcadfs(0); // root = 0

  for (int i = 0; i < MAXLOG - 1; i++) {
    for (int j = 0; j < n; j++) {
      if (parent[i][j] != -1) {
        parent[i + 1][j] = parent[i][parent[i][j]];

        if (parent[i + 1][j] != -1)
          min_dis[i + 1][j] = min(min_dis[i][j], min_dis[i][parent[i][j]]);
      }
    }
  }

  for (int i = 0; i < q; i++) {
    scanf("%d%d", &a, &b);
    a--;
    b--;

    int ans = INF;

    if (depth[a] < depth[b])
      swap(a, b);

    int d = depth[a] - depth[b];

    for (int j = 0; d > 0; j++) {
      if (d % 2) {
        ans = min(ans, min_dis[j][a]);
        a = parent[j][a];
      }
      d /= 2;
    }

    if (a != b) {
      for (int j = MAXLOG - 1; j >= 0; j--) {

        if (parent[j][a] != parent[j][b]) {
          ans = min(ans, min_dis[j][a]);
          ans = min(ans, min_dis[j][b]);
          a = parent[j][a];
          b = parent[j][b];
        }
      }

      ans = min(ans, min_dis[0][a]);
      ans = min(ans, min_dis[0][b]);
      a = parent[0][a];
      b = parent[0][b];
    }

    assert(a == b && a != -1);

    printf("%d\\n", ans);
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

typedef long long ll;

using namespace std;

struct TOWN {
  vector<pair<int, int>> road;
  int df;
  bool kaku;

  TOWN() : df(2000000000), kaku(false) {}

} town[100001];

int main() {
  int n, m, k, q, a, b, c;

  priority_queue<pair<int, int>> dij;

  scanf("%d%d%d%d", &n, &m, &k, &q);

  for (int i = 0; i < m; i++) {
    scanf("%d%d%d", &a, &b, &c);
    town[a].road.push_back(make_pair(b, -c));
    town[b].road.push_back(make_pair(a, -c));
  }

  for (int i = 0; i < k; i++) {
    scanf("%d", &a);
    town[a].df = 0;
    dij.push(make_pair(0, a));
  }

  while (!dij.empty()) {
    pair<int, int> p;

    p = dij.top();
    dij.pop();

    if (!town[p.second].kaku) {
      town[p.second].kaku = true;
      town[p.second].df = -p.first;

      for (int i = 0; i < town[p.second].road.size(); i++) {
        if (!town[town[p.second].road[i].first].kaku) {
          dij.push(make_pair(p.first + town[p.second].road[i].second,
                             town[p.second].road[i].first));
        }
      }
    }
  }

  for (int i = 0; i < q; i++) {

    for (int j = 1; j <= n; j++) {
      town[j].kaku = false;
    }

    scanf("%d%d", &a, &b);
    dij.push(make_pair(town[a].df, a));

    while (!dij.empty()) {
      pair<int, int> p;

      p = dij.top();
      dij.pop();

      if (p.second == b) {
        printf("%d\\n", p.first);
        break;
      }

      if (!town[p.second].kaku) {
        town[p.second].kaku = true;

        for (int i = 0; i < town[p.second].road.size(); i++) {
          if (!town[town[p.second].road[i].first].kaku) {
            dij.push(
                make_pair(min(p.first, town[town[p.second].road[i].first].df),
                          town[p.second].road[i].first));
          }
        }
      }
    }

    while (!dij.empty())
      dij.pop();
  }

  return 0;
}