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

#include <bits/stdc++.h>

#define int long long
#define double long double
#define ff first
#define ss second
#define endl '\\n'
#define ii pair<int, int>
#define mp make_pair
#define mt make_tuple
#define DESYNC                                                                 \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define pb push_back
#define vi vector<int>
#define vii vector<ii>
#define all(x) x.begin(), x.end()
#define EPS 1e-9
#define INF 1e18
#define ROOT 1
#define M 1000000007
const double PI = acos(-1);

using namespace std;

inline int mod(int n, int m) {
  int ret = n % m;
  if (ret < 0)
    ret += m;
  return ret;
}

int gcd(int a, int b) {
  if (a == 0)
    return b;
  else
    return gcd(b % a, a);
}

void solution() {
  int n;
  cin >> n;
  vector<vector<int>> dist(n, vector<int>(n));
  vector<vector<int>> mat(n, vector<int>(n));
  ii pos[n * n];
  int val = 0;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      mat[i][j] = val;
      pos[val++] = ii(i, j);
      dist[i][j] = min(i, n - 1 - i);
      dist[i][j] = min(dist[i][j], min(j, n - 1 - j));
    }
  }

  function<bool(int &, int &)> valid = [&](int &a, int &b) {
    return a >= 0 && b >= 0 && a < n && b < n;
  };

  function<void(ii)> update = [&](const ii x) {
    int dx[] = {1, 0, -1, 0};
    int dy[] = {0, 1, 0, -1};

    for (int i = 0; i < 4; i++) {
      int nx = x.ff + dx[i];
      int ny = x.ss + dy[i];
      if (!valid(nx, ny))
        continue;
      if (dist[x.ff][x.ss] + (mat[x.ff][x.ss] != -1) < dist[nx][ny]) {
        dist[nx][ny] = dist[x.ff][x.ss] + (mat[x.ff][x.ss] != -1);
        update(ii(nx, ny));
      }
    }
  };

  int ans = 0;
  for (int i = 0; i < n * n; i++) {
    int x;
    cin >> x;
    x--;
    ans += dist[pos[x].ff][pos[x].ss];
    mat[pos[x].ff][pos[x].ss] = -1;
    update(pos[x]);
  }
  cout << ans << endl;
}

int32_t main() {
  DESYNC;
  solution();
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

#define int long long
#define double long double
#define ff first
#define ss second
#define endl '\\n'
#define ii pair<int, int>
#define mp make_pair
#define mt make_tuple
#define DESYNC                                                                 \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
#define pb push_back
#define vi vector<int>
#define vii vector<ii>
#define all(x) x.begin(), x.end()
#define EPS 1e-9
#define INF 1e18
#define ROOT 1
#define M 1000000007
const double PI = acos(-1);

using namespace std;

inline int mod(int n, int m) {
  int ret = n % m;
  if (ret < 0)
    ret += m;
  return ret;
}

int gcd(int a, int b) {
  if (a == 0)
    return b;
  else
    return gcd(b % a, a);
}

ii uf[505][505];
int sz[505][505];
bool ok[505][505];

ii find(int x, int y) {
  if (uf[x][y] == ii(x, y))
    return ii(x, y);
  else
    return uf[x][y] = find(uf[x][y].ff, uf[x][y].ss);
}

bool merge(int a, int b, int c, int d) {
  ii x = find(a, b);
  ii y = find(c, d);

  if (x == y)
    return false;
  if (sz[x.ff][x.ss] < sz[y.ff][y.ss])
    swap(x, y);
  uf[y.ff][y.ss] = uf[x.ff][x.ss];
  sz[x.ff][x.ss] += sz[y.ff][y.ss];
  ok[x.ff][x.ss] |= ok[y.ff][y.ss];
  return true;
}

void solution() {
  int n;
  cin >> n;
  int mat[n + 2][n + 2];
  int dst[n + 2][n + 2];
  bool zero_path[n + 2][n + 2];
  vector<ii> pos(n * n + 1);
  int val = 1;
  for (int i = 1; i <= n; i++) {
    for (int j = 1; j <= n; j++) {
      pos[val] = ii(i, j);
      mat[i][j] = val++;
      zero_path[i][j] = false;
    }
  }

  for (int i = 0; i <= n + 1; i++) {
    mat[0][i] = mat[i][0] = mat[i][n + 1] = mat[n + 1][i] = -1;
    zero_path[0][i] = zero_path[i][0] = zero_path[i][n + 1] =
        zero_path[n + 1][i] = true;
  }

  for (int i = 0; i <= n + 1; i++) {
    for (int j = 0; j <= n + 1; j++) {
      dst[i][j] = INF;
      uf[i][j] = ii(i, j);
      sz[i][j] = 1;
      ok[i][j] = zero_path[i][j];
    }
  }

  int ans = 0;
  for (int i = 0; i < n * n; i++) {
    int st;
    cin >> st;
    int x = pos[st].ff;
    int y = pos[st].ss;
    vector<ii> vis;
    deque<ii> q;
    q.push_back(ii(x, y));
    vis.pb(ii(x, y));
    dst[x][y] = 0;
    while (q.size() > 0) {
      int cx = q.front().ff;
      int cy = q.front().ss;
      q.pop_front();

      ii grp = find(cx, cy);
      if (ok[grp.ff][grp.ss]) {
        ans += dst[cx][cy];
        break;
      }

      if (cx == 0 || cx == n + 1 || cy == 0 || cy == n + 1) {
        ans += dst[cx][cy];
        break;
      }

      // move
      if (dst[cx][cy] + (mat[cx + 1][cy] != -1) < dst[cx + 1][cy]) {
        dst[cx + 1][cy] = dst[cx][cy] + (mat[cx + 1][cy] != -1);
        vis.pb(ii(cx + 1, cy));
        if (mat[cx + 1][cy] == -1)
          q.push_front(ii(cx + 1, cy));
        else
          q.push_back(ii(cx + 1, cy));
      }

      if (dst[cx][cy] + (mat[cx - 1][cy] != -1) < dst[cx - 1][cy]) {
        dst[cx - 1][cy] = dst[cx][cy] + (mat[cx - 1][cy] != -1);
        vis.pb(ii(cx - 1, cy));
        if (mat[cx - 1][cy] == -1)
          q.push_front(ii(cx - 1, cy));
        else
          q.push_back(ii(cx - 1, cy));
      }

      if (dst[cx][cy] + (mat[cx][cy + 1] != -1) < dst[cx][cy + 1]) {
        dst[cx][cy + 1] = dst[cx][cy] + (mat[cx][cy + 1] != -1);
        vis.pb(ii(cx, cy + 1));
        if (mat[cx][cy + 1] == -1)
          q.push_front(ii(cx, cy + 1));
        else
          q.push_back(ii(cx, cy + 1));
      }

      if (dst[cx][cy] + (mat[cx][cy - 1] != -1) < dst[cx][cy - 1]) {
        dst[cx][cy - 1] = dst[cx][cy] + (mat[cx][cy - 1] != -1);
        vis.pb(ii(cx, cy - 1));
        if (mat[cx][cy - 1] == -1)
          q.push_front(ii(cx, cy - 1));
        else
          q.push_back(ii(cx, cy - 1));
      }
    }
    for (ii p : vis)
      dst[p.ff][p.ss] = INF;
    mat[x][y] = -1;
    if (mat[x - 1][y] == -1)
      merge(x - 1, y, x, y);
    if (mat[x + 1][y] == -1)
      merge(x + 1, y, x, y);
    if (mat[x][y - 1] == -1)
      merge(x, y - 1, x, y);
    if (mat[x][y + 1] == -1)
      merge(x, y + 1, x, y);
  }
  cout << ans << endl;
}

int32_t main() {
  DESYNC;
  solution();
}