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

#include <bits/stdc++.h>
using namespace std;

int h, w, k, xa, ya, xb, yb;
set<int> row[1000005], col[1000005];
vector<pair<int, int>> visited;
vector<vector<int>> dist, dp1, dp2;
vector<string> grid;

int main() {
  // ios_base::sync_with_stdio(false);
  // cin.tie(0);
  // cout.tie(0);

  cin >> h >> w >> k >> xa >> ya >> xb >> yb;
  xa--;
  ya--;
  xb--;
  yb--;
  grid.resize(h + 5);
  dist.assign(h + 5, vector<int>(w + 5, -1));
  dp1.assign(h + 5, vector<int>(w + 5, 0));
  dp2.assign(h + 5, vector<int>(w + 5, 0));
  for (int i = 0; i < h; i++) {
    cin >> grid[i];
    for (int j = 0; j < w; j++) {
      if (grid[i][j] == '.') {
        row[i].insert(j);
        col[j].insert(i);
      }
    }
  }

  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      if (j) {
        dp1[i][j] = dp1[i][j - 1];
      }

      if (grid[i][j] == '@') {
        dp1[i][j]++;
      }
    }
  }

  for (int j = 0; j < w; j++) {
    for (int i = 0; i < h; i++) {
      if (i) {
        dp2[i][j] = dp2[i - 1][j];
      }

      if (grid[i][j] == '@') {
        dp2[i][j]++;
      }
    }
  }

  queue<pair<int, int>> q;
  q.push({xa, ya});
  dist[xa][ya] = 0;
  row[xa].erase(ya);
  col[ya].erase(xa);

  while (!q.empty() && dist[xb][yb] == -1) {
    pair<int, int> u = q.front();
    q.pop();

    // left
    auto p = row[u.first].lower_bound(u.second);
    if (p != row[u.first].begin()) {
      p--;
      for (auto it = p; *it >= u.second - k; it--) {
        int res =
            dp1[u.first][u.second] - (*it == 0 ? 0 : dp1[u.first][*it - 1]);
        if (res != 0) {
          break;
        }

        dist[u.first][*it] = dist[u.first][u.second] + 1;
        q.push({u.first, *it});
        visited.push_back({u.first, *it});

        if (it == row[u.first].begin()) {
          break;
        }
      }
    }

    // right
    p = row[u.first].lower_bound(u.second);
    for (auto it = p; it != row[u.first].end() && *it <= u.second + k; it++) {
      int res =
          dp1[u.first][*it] - (u.second == 0 ? 0 : dp1[u.first][u.second - 1]);
      if (res != 0) {
        break;
      }

      dist[u.first][*it] = dist[u.first][u.second] + 1;
      q.push({u.first, *it});
      visited.push_back({u.first, *it});
    }

    // up
    p = col[u.second].lower_bound(u.first);
    if (p != col[u.second].begin()) {
      p--;

      for (auto it = p; *it >= u.first - k; it--) {
        int res =
            dp2[u.first][u.second] - (*it == 0 ? 0 : dp2[*it - 1][u.second]);
        if (res != 0) {
          break;
        }

        dist[*it][u.second] = dist[u.first][u.second] + 1;
        q.push({*it, u.second});
        visited.push_back({*it, u.second});

        if (it == col[u.second].begin()) {
          break;
        }
      }
    }

    // down
    p = col[u.second].lower_bound(u.first);
    for (auto it = p; it != col[u.second].end() && *it <= u.first + k; it++) {
      int res =
          dp2[*it][u.second] - (u.first == 0 ? 0 : dp2[u.first - 1][u.second]);
      if (res != 0) {
        break;
      }

      dist[*it][u.second] = dist[u.first][u.second] + 1;
      q.push({*it, u.second});
      visited.push_back({*it, u.second});
    }

    for (int i = 0; i < visited.size(); i++) {
      pair<int, int> v = visited[i];
      row[v.first].erase(v.second);
      col[v.second].erase(v.first);
    }
    visited.clear();
  }

  cout << dist[xb][yb] << '\\n';

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

int h, w, k, xa, ya, xb, yb, res;
set<int> row[1000005], col[1000005];
vector<pair<int, int>> visited;
vector<vector<int>> dist, dp1, dp2;
vector<string> grid;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  cin >> h >> w >> k >> xa >> ya >> xb >> yb;
  xa--;
  ya--;
  xb--;
  yb--;
  grid.resize(h + 5);
  dist.assign(h + 5, vector<int>(w + 5, -1));
  dp1.assign(h + 5, vector<int>(w + 5, 0));
  dp2.assign(h + 5, vector<int>(w + 5, 0));
  for (int i = 0; i < h; i++) {
    cin >> grid[i];
    for (int j = 0; j < w; j++) {
      if (grid[i][j] == '.') {
        row[i].insert(j);
        col[j].insert(i);
      }
    }
  }

  for (int i = 0; i < h; i++) {
    for (int j = 0; j < w; j++) {
      if (j) {
        dp1[i][j] = dp1[i][j - 1];
      }

      if (grid[i][j] == '@') {
        dp1[i][j]++;
      }
    }
  }

  for (int j = 0; j < w; j++) {
    for (int i = 0; i < h; i++) {
      if (i) {
        dp2[i][j] = dp2[i - 1][j];
      }

      if (grid[i][j] == '@') {
        dp2[i][j]++;
      }
    }
  }

  queue<pair<int, int>> q;
  q.push({xa, ya});
  dist[xa][ya] = 0;
  row[xa].erase(ya);
  col[ya].erase(xa);

  while (!q.empty() && dist[xb][yb] == -1) {
    pair<int, int> u = q.front();
    q.pop();

    int target = max(0, u.second - k);
    auto p = row[u.first].lower_bound(target);

    for (auto it = p; it != row[u.first].end() && *it <= u.second + k; it++) {
      if (*it > u.second) {
        res = dp1[u.first][*it] -
              (u.second == 0 ? 0 : dp1[u.first][u.second - 1]);
      } else {
        res = dp1[u.first][u.second] - (*it == 0 ? 0 : dp1[u.first][*it - 1]);
      }

      if (res != 0) {
        continue;
      }

      dist[u.first][*it] = dist[u.first][u.second] + 1;
      q.push({u.first, *it});
      visited.push_back({u.first, *it});
    }

    target = max(0, u.first - k);
    p = col[u.second].lower_bound(target);
    for (auto it = p; it != col[u.second].end() && *it <= u.first + k; it++) {
      if (*it > u.first) {
        res = dp2[*it][u.second] -
              (u.first == 0 ? 0 : dp2[u.first - 1][u.second]);
      } else {
        res = dp2[u.first][u.second] - (*it == 0 ? 0 : dp2[*it - 1][u.second]);
      }

      if (res != 0) {
        continue;
      }

      dist[*it][u.second] = dist[u.first][u.second] + 1;
      q.push({*it, u.second});
      visited.push_back({*it, u.second});
    }

    for (int i = 0; i < visited.size(); i++) {
      pair<int, int> v = visited[i];
      row[v.first].erase(v.second);
      col[v.second].erase(v.first);
    }
    visited.clear();
  }

  cout << dist[xb][yb] << '\\n';

  return 0;
}