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

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using ii = pair<int, int>;

const ll oo{1LL << 62};

const map<char, ii> dirs{
    {'N', ii(-1, 0)}, {'S', ii(1, 0)}, {'W', ii(0, -1)}, {'E', ii(0, 1)}};

ll solve(int H, int W, int K, int x1, int y1, int x2, int y2,
         vector<vector<ll>> &xs, vector<set<int>> &rs, vector<set<int>> &cs) {
  queue<ii> q;
  q.push(ii(x2, y2));

  rs[x2].erase(y2);
  cs[y2].erase(x2);

  xs[x2][y2] = 0;

  while (not q.empty()) {
    auto p = q.front();
    q.pop();

    auto x = p.first, y = p.second;
    auto dist = xs[x][y] + 1;

    vector<ii> ps;

    // East
    for (auto it = rs[x].lower_bound(y);
         it != rs[x].end() and (*it - y) <= K and xs[x][*it] != oo; ++it) {
      xs[x][*it] = dist;
      ps.push_back(ii(x, *it));
    }

    while (not ps.empty()) {
      auto p = ps.back();
      ps.pop_back();

      auto u = p.first, v = p.second;
      rs[u].erase(v);
      cs[v].erase(u);

      q.push(p);
    }

    // West
    for (auto it = rs[x].lower_bound(y);
         it != rs[x].begin() and (y - *prev(it)) <= K and
         xs[x][*prev(it)] != oo;
         --it) {
      xs[x][*prev(it)] = dist;
      ps.push_back(ii(x, *prev(it)));
    }

    while (not ps.empty()) {
      auto p = ps.back();
      ps.pop_back();

      auto u = p.first, v = p.second;
      rs[u].erase(v);
      cs[v].erase(u);

      q.push(p);
    }

    // South
    for (auto it = cs[y].lower_bound(x);
         it != cs[y].end() and (*it - x) <= K and xs[*it][y] != oo; ++it) {
      xs[*it][y] = dist;
      ps.push_back(ii(*it, y));
    }

    while (not ps.empty()) {
      auto p = ps.back();
      ps.pop_back();

      auto u = p.first, v = p.second;
      rs[u].erase(v);
      cs[v].erase(u);

      q.push(p);
    }

    // North
    for (auto it = cs[y].lower_bound(x);
         it != cs[y].begin() and (x - *prev(it)) <= K and
         xs[*prev(it)][y] != oo;
         --it) {
      xs[*prev(it)][y] = dist;
      ps.push_back(ii(*prev(it), y));
    }

    while (not ps.empty()) {
      auto p = ps.back();
      ps.pop_back();

      auto u = p.first, v = p.second;
      rs[u].erase(v);
      cs[v].erase(u);

      q.push(p);
    }

#ifdef DEBUG
    cout << "(" << x << ", " << y << "), xs = " << endl;
    for (int x = 1; x <= H; ++x)
      for (int y = 1; y <= W; ++y)
        cout << setw(3) << (xs[x][y] == oo ? -9 : xs[x][y])
             << (y == W ? '\\n' : ' ');
#endif
  }

  return xs[x1][y1];
}

int main() {
  ios::sync_with_stdio(false);

  int H, W, K;
  cin >> H >> W >> K;

  int x1, y1, x2, y2;
  cin >> x1 >> y1 >> x2 >> y2;

  vector<vector<ll>> xs(H + 2, vector<ll>(W + 2, -1));
  vector<set<int>> rs(H + 1), cs(W + 1);
  string line;

  for (int i = 1; i <= H; ++i)
    for (int j = 1; j <= W; ++j) {
      rs[i].insert(j);
      cs[j].insert(i);
    }

  for (int i = 1; i <= H; ++i) {
    cin >> line;

    for (int j = 1; j <= W; ++j) {
      if (line[j - 1] == '@') {
        xs[i][j] = oo;
      }
    }
  }

  auto ans = solve(H, W, K, x1, y1, x2, y2, xs, rs, cs);

  cout << ans << '\\n';

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using ii = pair<int, int>;

const ll oo{1LL << 62};

const vector<ii> dirs{ii(-1, 0), ii(1, 0), ii(0, -1), ii(0, 1)};

ll solve(int H, int W, int K, int x1, int y1, int x2, int y2,
         vector<vector<ll>> &xs, vector<int> &rs, vector<int> &cs) {
  queue<ii> q;
  q.push(ii(x2, y2));

  vector<vector<int>> found(H + 2, vector<int>(W + 2, 0));

  xs[x2][y2] = 0;
  found[x2][y2] = 1;
  ++rs[x2];
  ++cs[y2];

  while (not q.empty()) {
    auto p = q.front();
    q.pop();

    auto x = p.first, y = p.second;
    auto dist = xs[x][y] + 1;

    for (auto d : dirs) {
      auto dx = d.first, dy = d.second;
      auto u = x + dx, v = y + dy;

      while (u > 0 and u <= H and v > 0 and v <= W and abs(u - x) <= K and
             abs(v - y) <= K and xs[u][v] != oo) {
        if (not found[u][v]) {
          found[u][v] = 1;
          ++rs[u];
          ++cs[v];
          q.push(ii(u, v));
          xs[u][v] = dist;
        } else {
          xs[u][v] = xs[u][v] == -1 ? dist : min(dist, xs[u][v]);

          if ((dx and rs[u] == H) or (dy and cs[v] == W))
            break;
        }

        u += dx;
        v += dy;
      }
    }
#ifdef DEBUG
    cout << "(" << x << ", " << y << "), xs = " << endl;
    for (int x = 1; x <= H; ++x)
      for (int y = 1; y <= W; ++y)
        cout << setw(3) << (xs[x][y] == oo ? -9 : xs[x][y])
             << (y == W ? '\\n' : ' ');
#endif
  }

#ifdef DEBUG
  cout << "--- x1 = " << x1 << ", y1 = " << y1 << ", xs = " << xs[x1][y1]
       << endl;
#endif

  return xs[x1][y1];
}

int main() {
  ios::sync_with_stdio(false);

  int H, W, K;
  cin >> H >> W >> K;

  int x1, y1, x2, y2;
  cin >> x1 >> y1 >> x2 >> y2;

  vector<vector<ll>> xs(H + 2, vector<ll>(W + 2, -1));
  vector<int> rs(H + 2, 0), cs(W + 2, 0);

  for (int i = 1; i <= H; ++i) {
    string s;
    cin >> s;

    for (int j = 1; j <= W; ++j) {
      if (s[j - 1] == '@') {
        ++rs[i];
        ++cs[j];
        xs[i][j] = oo;
      }
    }
  }

  auto ans = solve(H, W, K, x1, y1, x2, y2, xs, rs, cs);

  cout << ans << '\\n';

  return 0;
}