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

#include <bits/stdc++.h>
using namespace std;
int const nax = 1e6 + 10;
char s[nax];
pair<int, int> q[nax];

int main() {
  vector<vector<int>> a;
  vector<vector<int>> dist;
  vector<vector<int>> vis;
  int n, m, k;
  scanf("%d %d %d", &n, &m, &k);
  a.resize(n);

  pair<int, int> st, ed;

  scanf("%d %d", &st.first, &st.second);
  scanf("%d %d", &ed.first, &ed.second);
  --st.first, --st.second;
  --ed.first, --ed.second;

  for (int i = 0; i < n; ++i) {
    scanf("%s", s);
    for (int j = 0; j < m; ++j) {
      a[i].push_back(s[j] == '@' ? 0 : 1);
    }
  }

  vis.resize(n);
  dist.resize(n);
  for (int i = 0; i < n; ++i) {
    vis[i].resize(m, false);
    dist[i].resize(m, 0);
  }

  set<int> rpos[n];
  set<int> cpos[m];

  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < m; ++j)
      rpos[i].insert(j);
  }

  for (int j = 0; j < m; ++j)
    for (int i = 0; i < n; ++i) {
      cpos[j].insert(i);
    }
  fprintf(stderr, "ALIVE\\n");

  constexpr pair<int, int> invalid_pair = make_pair(-1, -1);

  auto get_neighbour = [&](int i, int j, int dir) {
    if (dir == 0) {
      auto it = cpos[j].lower_bound(i);
      if (cpos[j].empty() || it == cpos[j].begin())
        return invalid_pair;
      --it;
      if (*it >= i)
        return invalid_pair;
      return make_pair(*it, j);
    } else if (dir == 1) {
      auto it = cpos[j].upper_bound(i);
      if (it == cpos[j].end())
        return invalid_pair;
      return make_pair(*it, j);
    } else if (dir == 2) {
      auto it = rpos[i].lower_bound(j);
      if (rpos[i].empty() || it == rpos[i].begin())
        return invalid_pair;
      --it;
      if (*it >= j)
        return invalid_pair;
      return make_pair(i, *it);
    } else if (dir == 3) {
      auto it = rpos[i].upper_bound(j);
      if (it == rpos[i].end())
        return invalid_pair;
      return make_pair(i, *it);
    } else
      throw;
  };

  auto erase_from_set = [&](int i, int j) -> void {
    if (cpos[j].find(i) != cpos[j].end())
      cpos[j].erase(cpos[j].find(i));

    if (rpos[i].find(j) != rpos[i].end())
      rpos[i].erase(rpos[i].find(j));
  };

  fprintf(stderr, "ALIVE2\\n");

  const int inf = 1e9;
  for (int i = 0; i < n; ++i)
    for (int j = 0; j < m; ++j) {
      dist[i][j] = inf;
    }

  dist[st.first][st.second] = 0;

  int topQ = 0;

  q[topQ++] = st;
  erase_from_set(st.first, st.second);

  for (int g = 0; g < topQ; ++g) {

    int x = q[g].first;
    int y = q[g].second;
    int cur_dist = dist[x][y];

    vector<pair<int, int>> to_push;

    for (int dir = 0; dir < 4; ++dir) {
      pair<int, int> cur = make_pair(x, y);
      while (true) {
        auto u = get_neighbour(cur.first, cur.second, dir);
        if (u == invalid_pair)
          break;
        if (!a[u.first][u.second])
          break;
        if (abs(u.first - x) + abs(u.second - y) > k)
          break;
        if (u == cur)
          break;
        cur = u;
        to_push.push_back(u);
      }
    }

    for (auto &it : to_push) {
      dist[it.first][it.second] = cur_dist + 1;
      erase_from_set(it.first, it.second);
      q[topQ++] = it;
    }
  }

  /* for (int i = 0 ; i < n;  ++ i, printf("\\n"))
       for (int j = 0 ; j < m ; ++ j)
           printf("%d\\t", dist[i][j]);*/

  printf("%d\\n",
         dist[ed.first][ed.second] == inf ? -1 : dist[ed.first][ed.second]);
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
int const nax = 1e6 + 10;
char s[nax];
pair<int, int> q[nax];

int main() {
  vector<vector<int>> a;
  vector<vector<int>> dist;
  vector<vector<int>> vis;
  int n, m, k;
  scanf("%d %d %d", &n, &m, &k);
  a.resize(n);

  pair<int, int> st, ed;

  scanf("%d %d", &st.first, &st.second);
  scanf("%d %d", &ed.first, &ed.second);
  --st.first, --st.second;
  --ed.first, --ed.second;

  for (int i = 0; i < n; ++i) {
    scanf("%s", s);
    for (int j = 0; j < m; ++j) {
      a[i].push_back(s[j] == '@' ? 0 : 1);
    }
  }

  vector<vector<int>> up;
  vector<vector<int>> down;
  vector<vector<int>> left;
  vector<vector<int>> right;

  auto resize_vec = [&](vector<vector<int>> &f) -> void {
    f.resize(n);
    for (int i = 0; i < n; ++i) {
      f[i].resize(m, -1);
    }
  };

  resize_vec(up);
  resize_vec(down);
  resize_vec(left);
  resize_vec(right);
  resize_vec(dist);

  vis.resize(n);
  for (int i = 0; i < n; ++i) {
    vis[i].resize(m, false);
  }

  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < m; ++j) {
      if (!a[i][j])
        continue;
      if (i > 0 && a[i - 1][j]) {
        up[i][j] = i - 1;
      }

      if (i + 1 < n && a[i + 1][j]) {
        down[i][j] = i + 1;
      }

      if (j > 0 && a[i][j - 1])
        left[i][j] = j - 1;

      if (j + 1 < m && a[i][j + 1])
        right[i][j] = j + 1;
    }
  }

  auto get_neightbor = [&](pair<int, int> p, int dir) {
    if (dir == 0)
      return make_pair(up[p.first][p.second], p.second);
    if (dir == 1)
      return make_pair(down[p.first][p.second], p.second);
    if (dir == 2)
      return make_pair(p.first, left[p.first][p.second]);
    return make_pair(p.first, right[p.first][p.second]);
  };

  const int inf = 1e9;
  for (int i = 0; i < n; ++i)
    for (int j = 0; j < m; ++j) {
      dist[i][j] = inf;
    }

  dist[st.first][st.second] = 0;

  int topQ = 0;

  if (a[st.first][st.second]) {
    q[topQ++] = st;
    {
      int i = st.first;
      int j = st.second;
      if (up[i][j] != -1) {
        down[up[i][j]][j] = down[i][j];
      }
      if (down[i][j] != n && down[i][j] != -1) {
        up[down[i][j]][j] = up[i][j];
      }
      if (left[i][j] != -1) {
        right[i][left[i][j]] = right[i][j];
      }
      if (right[i][j] != m && right[i][j] != -1) {
        left[i][right[i][j]] = left[i][j];
      }
    }
  }
  vis[st.first][st.second] = true;

  function<void(int, int)> dfs = [&](int i, int j) {
    vis[i][j] = true;
    if (i > 0 && a[i - 1][j] && !vis[i - 1][j]) {
      dfs(i - 1, j);
    }

    if (i + 1 < n && a[i + 1][j] && !vis[i + 1][j]) {
      dfs(i + 1, j);
    }

    if (j > 0 && a[i][j - 1] && !vis[i][j - 1]) {
      dfs(i, j - 1);
    }
    if (j + 1 < m && a[i][j + 1] && !vis[i][j + 1])
      dfs(i, j + 1);
  };

  dfs(st.first, st.second);
  if (!vis[ed.first][ed.second])
    return !printf("-1");

  for (int i = 0; i < n; ++i)
    for (int j = 0; j < m; ++j)
      vis[i][j] = false;
  vis[st.first][st.second] = true;

  for (int g = 0; g < topQ; ++g) {

    int x = q[g].first;
    int y = q[g].second;
    int cur_dist = dist[x][y];

    if (make_pair(x, y) == ed) {
      return !printf("%d\\n", dist[x][y]);
    }

    vector<pair<int, int>> to_push;

    for (int dir = 0; dir < 4; ++dir) {
      int cur_x = x, cur_y = y;
      while (true) {
        auto u = get_neightbor(make_pair(cur_x, cur_y), dir);
        if (u.first == -1 || u.second == -1)
          break;

        if (u.first == cur_x && u.second == cur_y)
          break;

        if (abs(u.first - x) + abs(u.second - y) > k)
          break;
        if (!vis[u.first][u.second]) {
          to_push.emplace_back(u.first, u.second);
          q[topQ++] = u;
          vis[u.first][u.second] = true;
        }
        cur_x = u.first;
        cur_y = u.second;
      }
    }

    for (auto &p : to_push) {
      int i = p.first;
      int j = p.second;

      dist[i][j] = min(dist[i][j], cur_dist + 1);
      if (up[i][j] != -1) {
        down[up[i][j]][j] = down[i][j];
      }
      if (down[i][j] != n && down[i][j] != -1) {
        up[down[i][j]][j] = up[i][j];
      }
      if (left[i][j] != -1) {
        right[i][left[i][j]] = right[i][j];
      }
      if (right[i][j] != m && right[i][j] != -1) {
        left[i][right[i][j]] = left[i][j];
      }
    }

    for (auto &p : to_push) {
      if (!vis[p.first][p.second])
        q[topQ++] = p;
      vis[p.first][p.second] = true;
    }
  }

  if (!vis[ed.first][ed.second]) {
    return !printf("-1");
  }

  printf("%d\\n", dist[ed.first][ed.second]);
}