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

#include <algorithm>
#include <fstream>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <vector>

using namespace std;

const int DIV = 998244353;

int main(int argc, char *argv[]) {
#ifdef VLAD_LOCAL
  ifstream in("in.txt");
#else
  istream &in = cin;
#endif
  int h, w, k;
  in >> h >> w >> k;
  int x1, y1, x2, y2;
  in >> x1 >> y1 >> x2 >> y2;
  x1--;
  y1--;
  x2--;
  y2--;
  vector<string> c(h);
  for (int i = 0; i < h; i++)
    in >> c[i];

  vector<set<int>> waitX(h);
  vector<set<int>> waitY(w);
  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      if (c[i][j] == '.') {
        waitX[i].insert(j);
        waitY[j].insert(i);
      }
    }
  }

  vector<vector<int>> cntX(h);
  for (int i = 0; i < h; i++)
    cntX[i].resize(w + 1, 0);
  vector<vector<int>> cntY(w);
  for (int i = 0; i < w; i++)
    cntY[i].resize(h + 1, 0);
  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      bool has = c[i][j] == '@';
      cntX[i][j + 1] = cntX[i][j] + has;
      cntY[j][i + 1] = cntY[j][i] + has;
    }
  }

  vector<vector<int>> d(h);
  for (int i = 0; i < h; i++)
    d[i].resize(w, -1);
  d[x1][y1] = 0;
  waitX[x1].erase(y1);
  waitY[y1].erase(x1);
  priority_queue<pair<int, pair<int, int>>> q;
  q.emplace(0, make_pair(x1, y1));
  while (!q.empty()) {
    pair<int, int> v = q.top().second;
    int curD = -q.top().first;
    q.pop();

    vector<pair<int, int>> toDel;
    set<int>::iterator itV, it;

    itV = waitX[v.first].lower_bound(v.second);
    for (it = itV; it != waitX[v.first].end(); it++) {
      int cnt = cntX[v.first][(*it) + 1] - cntX[v.first][v.second];
      int dist = (*it) - v.second;
      if (cnt || dist > k)
        break;
      int x = v.first;
      int y = *it;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }
    for (it = itV; it != waitX[v.first].begin();) {
      it--;
      int cnt = cntX[v.first][v.second + 1] - cntX[v.first][(*it)];
      int dist = -(*it) + v.second;
      if (cnt || dist > k)
        break;
      int x = v.first;
      int y = *it;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }

    itV = waitY[v.second].lower_bound(v.first);
    for (it = itV; it != waitY[v.second].end(); it++) {
      int cnt = cntY[v.second][(*it) + 1] - cntY[v.second][v.first];
      int dist = (*it) - v.first;
      if (cnt || dist > k)
        break;
      int x = *it;
      int y = v.second;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }
    for (it = itV; it != waitY[v.second].begin();) {
      it--;
      int cnt = cntY[v.second][v.first + 1] - cntY[v.second][(*it)];
      int dist = -(*it) + v.first;
      if (cnt || dist > k)
        break;
      int x = *it;
      int y = v.second;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }

    for (int i = 0; i < toDel.size(); i++) {
      waitX[toDel[i].first].erase(toDel[i].second);
      waitY[toDel[i].second].erase(toDel[i].first);
    }
  }

  printf("%d\\n", d[x2][y2]);

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <fstream>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <vector>

using namespace std;

const int DIV = 998244353;

int main(int argc, char *argv[]) {
#ifdef VLAD_LOCAL
  ifstream in("in.txt");
#else
  istream &in = cin;
#endif
  int h, w, k;
  in >> h >> w >> k;
  int x1, y1, x2, y2;
  in >> x1 >> y1 >> x2 >> y2;
  x1--;
  y1--;
  x2--;
  y2--;
  vector<string> c(h);
  for (int i = 0; i < h; i++)
    in >> c[i];

  vector<set<int>> waitX(h);
  vector<set<int>> waitY(w);
  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      if (c[i][j] == '.') {
        waitX[i].insert(j);
        waitY[j].insert(i);
      }
    }
  }

  vector<vector<int>> cntX(h);
  for (int i = 0; i < h; i++)
    cntX[i].resize(w + 1, 0);
  vector<vector<int>> cntY(w);
  for (int i = 0; i < w; i++)
    cntY[i].resize(h + 1, 0);
  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      bool has = c[i][j] == '@';
      cntX[i][j + 1] = cntX[i][j] + has;
      cntY[j][i + 1] = cntY[j][i] + has;
    }
  }

  vector<vector<int>> d(h);
  for (int i = 0; i < h; i++)
    d[i].resize(w, -1);
  d[x1][y1] = 0;
  waitX[x1].erase(y1);
  waitY[y1].erase(x1);
  priority_queue<pair<int, pair<int, int>>> q;
  q.emplace(0, make_pair(x1, y1));
  while (!q.empty()) {
    pair<int, int> v = q.top().second;
    int curD = -q.top().first;
    q.pop();

    vector<pair<int, int>> toDel;
    set<int>::iterator itV, it;

    itV = waitX[v.first].lower_bound(v.second);
    for (it = itV; it != waitX[v.first].end(); it++) {
      int cnt = cntX[v.first][(*it) + 1] - cntX[v.first][v.second];
      int dist = (*it) - v.second;
      if (cnt || dist > k)
        break;
      int x = v.first;
      int y = *it;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }
    for (it = itV; it != waitX[v.first].begin();) {
      it--;
      int cnt = cntX[v.first][v.second + 1] - cntX[v.first][(*it)];
      int dist = -(*it) + v.second;
      if (cnt || dist > k)
        break;
      int x = v.first;
      int y = *it;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }

    itV = waitY[v.second].lower_bound(v.first);
    for (it = itV; it != waitY[v.second].end(); it++) {
      int cnt = cntY[v.second][(*it) + 1] - cntY[v.second][v.first];
      int dist = (*it) - v.first;
      if (cnt || dist > k)
        break;
      int x = *it;
      int y = v.second;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }
    for (it = itV; it != waitY[v.second].begin();) {
      it--;
      int cnt = cntY[v.second][v.first + 1] - cntY[v.second][(*it)];
      int dist = -(*it) + v.first;
      if (cnt || dist > k)
        break;
      int x = *it;
      int y = v.second;
      if (d[x][y] < 0) {
        d[x][y] = curD + 1;
        q.emplace(-d[x][y], make_pair(x, y));
        toDel.emplace_back(x, y);
      }
    }
  }

  printf("%d\\n", d[x2][y2]);

  return 0;
}