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

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

mt19937 mt(736);

const int dx[] = {0, 0, -1, 1}, dy[] = {-1, 1, 0, 0};

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  int n;

  cin >> n;

  vector<int> ord(n * n);

  for (auto &it : ord) {
    cin >> it;
    it--;
  }

  vector<vector<bool>> alive(n, vector<bool>(n, true));
  vector<vector<int>> ans(n, vector<int>(n));

  for (int y = 0; y < ans.size(); y++)
    for (int x = 0; x < ans[y].size(); x++)
      ans[y][x] = min(min(y, (int)ans.size() - 1 - y),
                      min(x, (int)ans.size() - 1 - x)) +
                  1;

  int sum = 0;

  function<void(int, int, int)> dfs = [&ans, &alive, &dfs](int x, int y,
                                                           int ns) {
    if (y < 0 || y >= ans.size() || x < 0 || x >= ans[y].size())
      return;

    if (alive[y][x])
      ns++;

    if (ns >= ans[y][x])
      return;

    ans[y][x] = ns;

    for (int i = 0; i < 4; i++)
      dfs(x + dx[i], y + dy[i], ns);
  };

  for (auto it : ord) {
    int y = it / (int)ans.size();
    int x = it % (int)ans.size();

    auto add = ans[y][x] - 1;
    sum += add;
    alive[y][x] = false;

    dfs(x, y, add);
  }

  cout << sum << endl;
}

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

  cout << fixed;

#ifdef LOCAL
  auto st = clock();

  ifstream fin("../input.txt");

  solve(fin);

  cout << "clock: " << setprecision(4)
       << (clock() - st) / (double)CLOCKS_PER_SEC << endl;
#else
  solve();
#endif

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

mt19937 mt(736);

class dsu;

void update(dsu &white, vector<vector<int>> &ans, int x, int y, int ns);

class dsu {
  vector<int> par, val;
  vector<vector<int>> neigh;

public:
  dsu(int n = 0) : par(n, -1), val(n, n + 1), neigh(n) { val.back() = 0; }

  int get_class(int v) { return par[v] == -1 ? v : par[v] = get_class(par[v]); }

  bool unite(int a, int b) {
    a = get_class(a);
    b = get_class(b);

    if (a == b)
      return false;

    if (neigh[a].size() < neigh[b].size())
      swap(a, b);

    val[a] = min(val[a], val[b]);
    par[b] = a;
    neigh[a].insert(neigh[a].end(), neigh[b].begin(), neigh[b].end());
    neigh[b].clear();

    return true;
  }

  void add(int wh, int x) { neigh[get_class(wh)].push_back(x); }

  void start_update(vector<vector<int>> &ans, int wh) {
    wh = get_class(wh);

    for (int i = 0; i < neigh[wh].size(); i++) {
      auto &it = neigh[wh][i];

      if (ans[it / ans[0].size()][it % ans[0].size()] == -1) {
        swap(it, neigh[wh].back());
        neigh[wh].pop_back();
        i--;
      } else
        ::update(*this, ans, it % ans[0].size(), it / ans[0].size(), val[wh]);
    }
  }

  void update(vector<vector<int>> &ans, int wh, int ns) {
    wh = get_class(wh);

    if (ns < val[wh]) {
      val[wh] = ns;

      start_update(ans, wh);
    }
  }
};

const int dx[] = {-1, 1, 0, 0}, dy[] = {0, 0, -1, 1};

void update(dsu &white, vector<vector<int>> &ans, int x, int y, int ns) {
  if (ans[y][x] == -1) {
    white.update(ans, y * ans[0].size() + x, ns);

    return;
  }

  ns++;

  if (ns < ans[y][x]) {
    ans[y][x] = ns;

    for (int i = 0; i < 4; i++) {
      if (0 <= y + dy[i] && y + dy[i] < ans[0].size() && 0 <= x + dx[i] &&
          x + dx[i] < ans[0].size())
        update(white, ans, x + dx[i], y + dy[i], ns);
    }
  }
}

void solve(istream &cin = std::cin, ostream &cout = std::cout) {
  int n;

  cin >> n;

  vector<int> perm(n * n);

  for (auto &it : perm) {
    cin >> it;
    it--;
  }

  vector<vector<int>> ans(n, vector<int>(n));

  for (int y = 0; y < ans.size(); y++)
    for (int x = 0; x < ans[y].size(); x++)
      ans[y][x] = min(min(y, (int)ans.size() - 1 - y),
                      min(x, (int)ans.size() - 1 - x)) +
                  1;
  ans.emplace_back(1, -1);

  dsu white(n * n + 1);
  ll sum = 0;

  for (auto it : perm) {
    int y = it / n;
    int x = it % n;

    assert(ans[y][x] > 0);

    int ns = ans[y][x];

    //		cerr << it << ' ' << ns - 1 << endl;

    sum += ns - 1;
    ans[y][x] = -1;
    white.update(ans, it, ns - 1);

    for (int i = 0; i < 4; i++) {
      if (0 <= y + dy[i] && y + dy[i] < n && 0 <= x + dx[i] && x + dx[i] < n &&
          ans[y + dy[i]][x + dx[i]] != -1)
        white.add(it, (y + dy[i]) * n + x + dx[i]);
    }

    for (int i = 0; i < 4; i++) {
      if (0 <= y + dy[i] && y + dy[i] < n && 0 <= x + dx[i] && x + dx[i] < n) {
        if (ans[y + dy[i]][x + dx[i]] == -1)
          white.unite(it, (y + dy[i]) * n + x + dx[i]);
      } else
        white.unite(it, n * n);
    }

    white.start_update(ans, it);
  }

  cout << sum << endl;
}

void gen(ostream &cout = std::cout) {
  const int n = 4;

  cout << n << endl;

  vector<int> arr(n * n);
  iota(arr.begin(), arr.end(), 1);
  shuffle(arr.begin(), arr.end(), mt);

  for (auto it : arr)
    cout << it << ' ';
  cout << endl;
}

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

  cout << fixed;

#ifdef LOCAL
  auto st = clock();

  ifstream fin("../input.txt");

  solve(fin);

  while (!true) {
    stringstream ss;

    gen(ss);

    cout << ss.str() << endl;

    solve(ss);
  }

  cout << "clock: " << setprecision(4)
       << (clock() - st) / (double)CLOCKS_PER_SEC << endl;
#else
  solve();
#endif

  return 0;
}