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

#include <bits/stdc++.h>
#define int long long
#pragma GCC optimize("O3,Ofast,no-stack-protector,unroll-loops,fast-math")
using namespace std;

int H, W, K, a, b, c, d;
vector<vector<char>> kek;
vector<vector<int>> dp;

class UnionFind {
public:
  vector<int> to;
  vector<int> ba;
  vector<bool> keep;
  UnionFind(int n = 0) {
    to.resize(n + 2);
    ba.resize(n + 2);
    keep.resize(n + 2);
    for (int i = 1; i <= n; i++) {
      to[i] = i + 1;
      ba[i] = i - 1;
      keep[i] = true;
    }
    keep[0] = keep[n + 1] = true;
  }
  void cut(int i) { keep[i] = false; }
  int inc(int i) {
    if (keep[to[i]]) {
      return to[i];
    }
    return to[i] = inc(to[i]);
  }
  int dec(int i) {
    if (keep[ba[i]]) {
      return ba[i];
    }
    return ba[i] = dec(ba[i]);
  }
};

vector<UnionFind> rows;
vector<UnionFind> cols;

void kill(int x, int y) {
  rows[x].cut(y);
  cols[y].cut(x);
}

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

  cin >> H >> W >> K >> a >> b >> c >> d;
  kek.resize(H + 1);
  dp.resize(H + 1);
  for (int i = 1; i <= H; i++) {
    kek[i].resize(W + 1);
    dp[i].resize(W + 1, -1);
    for (int j = 1; j <= W; j++) {
      cin >> kek[i][j];
    }
  }
  rows.resize(H + 1);
  for (int i = 1; i <= H; i++) {
    rows[i] = UnionFind(W);
  }
  cols.resize(W + 1);
  for (int i = 1; i <= W; i++) {
    cols[i] = UnionFind(H);
  }
  queue<pair<int, int>> qp;
  qp.push({a, b});
  dp[a][b] = 0;
  kill(a, b);
  while (!qp.empty()) {
    pair<int, int> pr = qp.front();
    qp.pop();
    int x = pr.first;
    int y = pr.second;
    for (int i = x + 1; i <= min(H, x + K); i = cols[y].inc(i)) {
      if (kek[i][y] == '@')
        break;
      if (dp[i][y] != -1)
        continue;
      dp[i][y] = dp[x][y] + 1;
      kill(i, y);
      qp.push({i, y});
    }
    for (int i = x - 1; i >= max(1LL, x - K); i = cols[y].dec(i)) {
      if (kek[i][y] == '@')
        break;
      if (dp[i][y] != -1)
        continue;
      dp[i][y] = dp[x][y] + 1;
      kill(i, y);
      qp.push({i, y});
    }
    for (int i = y + 1; i <= min(W, y + K); i = rows[x].inc(i)) {
      if (kek[x][i] == '@')
        break;
      if (dp[x][i] != -1)
        continue;
      dp[x][i] = dp[x][y] + 1;
      kill(x, i);
      qp.push({x, i});
    }
    for (int i = y - 1; i >= max(1LL, y - K); i = rows[x].dec(i)) {
      if (kek[x][i] == '@')
        break;
      if (dp[x][i] != -1)
        continue;
      dp[x][i] = dp[x][y] + 1;
      kill(x, i);
      qp.push({x, i});
    }
  }
  cout << dp[c][d] << '\\n';

  fflush(stdout);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define int long long
#pragma GCC optimize("O3,Ofast,no-stack-protector,unroll-loops,fast-math")
using namespace std;

int H, W, K, a, b, c, d;
vector<vector<char>> kek;
vector<vector<int>> dp;

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

  cin >> H >> W >> K >> a >> b >> c >> d;
  kek.resize(H + 1);
  dp.resize(H + 1);
  for (int i = 1; i <= H; i++) {
    kek[i].resize(W + 1);
    dp[i].resize(W + 1, -1);
    for (int j = 1; j <= W; j++) {
      cin >> kek[i][j];
    }
  }
  queue<pair<int, int>> qp;
  qp.push({a, b});
  dp[a][b] = 0;
  while (!qp.empty()) {
    pair<int, int> pr = qp.front();
    qp.pop();
    int x = pr.first;
    int y = pr.second;
    for (int i = x + 1; i <= min(H, x + K); i++) {
      if (kek[i][y] == '@')
        break;
      if (dp[i][y] != -1)
        continue;
      dp[i][y] = dp[x][y] + 1;
      qp.push({i, y});
    }
    for (int i = x - 1; i >= max(1LL, x - K); i--) {
      if (kek[i][y] == '@')
        break;
      if (dp[i][y] != -1)
        continue;
      dp[i][y] = dp[x][y] + 1;
      qp.push({i, y});
    }
    for (int i = y + 1; i <= min(W, y + K); i++) {
      if (kek[x][i] == '@')
        break;
      if (dp[x][i] != -1)
        continue;
      dp[x][i] = dp[x][y] + 1;
      qp.push({x, i});
    }
    for (int i = y - 1; i >= max(1LL, y - K); i--) {
      if (kek[x][i] == '@')
        break;
      if (dp[x][i] != -1)
        continue;
      dp[x][i] = dp[x][y] + 1;
      qp.push({x, i});
    }
  }
  cout << dp[c][d] << '\\n';

  fflush(stdout);
  return 0;
}