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

#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
const int inf = 1e9;

vector<int> bfs(vector<vector<int>> &adj, int s) {
  vector<int> ret(adj.size(), inf);
  ret[s] = 0;
  queue<pair<int, int>> wait;
  wait.push(make_pair(s, 0));
  while (!wait.empty()) {
    int pos = wait.front().first;
    int cost = wait.front().second;
    wait.pop();
    for (int i = 0; i < (int)adj[pos].size(); i++) {
      if (ret[adj[pos][i]] == inf) {
        ret[adj[pos][i]] = cost + 1;
        wait.push(make_pair(adj[pos][i], cost + 1));
      }
    }
  }
  return ret;
}

int main() {
  int n, m, s, t;
  cin >> n >> m >> s >> t;
  vector<vector<int>> adj(n + 1);
  vector<int> x(m), y(m);
  for (int i = 0; i < m; i++) {
    cin >> x[i] >> y[i];
    adj[x[i]].push_back(y[i]);
    adj[y[i]].push_back(x[i]);
  }

  vector<int> froms = bfs(adj, s);
  vector<int> fromt = bfs(adj, t);
  int dist_st = froms[t];
  vector<long long int> bsums(n + 1, 0), bsumt(n + 1, 0);
  for (int i = 1; i < n + 1; i++) {
    if (froms[i] != inf)
      bsums[froms[i]]++;
    if (fromt[i] != inf)
      bsumt[fromt[i]]++;
  }

  long long int ans = 0;
  for (int i = 0; i < dist_st - 1; i++) {
    int j = dist_st - i - 2;
    if (i < j) {
      ans += bsums[i] * bsumt[j] + bsums[j] * bsumt[i];
    } else if (i == j) {
      ans += bsums[i] * bsumt[j];
    }
  }
  for (int i = 0; i < m; i++) {
    if (froms[x[i]] + fromt[y[i]] == dist_st - 2 ||
        froms[y[i]] + fromt[x[i]] == dist_st - 2) {
      ans--;
    }
  }
  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <queue>
#include <set>
#include <vector>
using namespace std;
const int inf = 1e9;

vector<int> bfs(vector<vector<int>> &adj, int s) {
  vector<int> ret(adj.size(), inf);
  ret[s] = 0;
  queue<pair<int, int>> wait;
  wait.push(make_pair(s, 0));
  while (!wait.empty()) {
    int pos = wait.front().first;
    int cost = wait.front().second;
    wait.pop();
    for (int i = 0; i < (int)adj[pos].size(); i++) {
      if (ret[adj[pos][i]] == inf) {
        ret[adj[pos][i]] = cost + 1;
        wait.push(make_pair(adj[pos][i], cost + 1));
      }
    }
  }
  return ret;
}

int main() {
  int n, m, s, t;
  cin >> n >> m >> s >> t;
  vector<vector<int>> adj(n + 1);
  set<pair<int, int>> elist;
  for (int i = 0; i < m; i++) {
    int x, y;
    cin >> x >> y;
    adj[x].push_back(y);
    adj[y].push_back(x);
    elist.insert(make_pair(min(x, y), max(x, y)));
  }

  vector<int> froms = bfs(adj, s);
  vector<int> fromt = bfs(adj, t);
  int dist_st = froms[t];
  vector<pair<int, int>> pfromt(n + 1);
  for (int i = 0; i < n + 1; i++) {
    pfromt[i] = make_pair(fromt[i], i);
  }
  sort(pfromt.begin(), pfromt.end());

  int ans = 0;
  for (int i = 1; i < n + 1; i++) {
    int target = dist_st - froms[i] - 2;
    vector<pair<int, int>>::iterator it;
    it = lower_bound(pfromt.begin(), pfromt.end(), make_pair(target, 0));
    while (it->first == target) {
      int vmin = min(i, it->second);
      int vmax = max(i, it->second);
      if (elist.find(make_pair(vmin, vmax)) == elist.end()) {
        ans++;
      }
      it++;
    }
  }
  cout << ans << endl;

  return 0;
}