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

#include <bits/stdc++.h>
// #include <atcoder/all>
// using namespace atcoder;
#define INF 2000000000000000000
#define ll long long
#define pll pair<ll, ll>
using namespace std;

// ===========Disjoint Set Union================
struct dsu {
public:
  dsu() : _n(0) {}
  dsu(int n) : _n(n), parent_or_size(n, -1) {}

  int merge(int a, int b) {
    assert(0 <= a && a < _n);
    assert(0 <= b && b < _n);
    int x = leader(a), y = leader(b);
    if (x == y)
      return x;
    if (-parent_or_size[x] < -parent_or_size[y])
      std::swap(x, y);
    parent_or_size[x] += parent_or_size[y];
    parent_or_size[y] = x;
    return x;
  }

  bool same(int a, int b) {
    assert(0 <= a && a < _n);
    assert(0 <= b && b < _n);
    return leader(a) == leader(b);
  }

  int leader(int a) {
    assert(0 <= a && a < _n);
    if (parent_or_size[a] < 0)
      return a;
    return parent_or_size[a] = leader(parent_or_size[a]);
  }

  int size(int a) {
    assert(0 <= a && a < _n);
    return -parent_or_size[leader(a)];
  }

  std::vector<std::vector<int>> groups() {
    std::vector<int> leader_buf(_n), group_size(_n);
    for (int i = 0; i < _n; i++) {
      leader_buf[i] = leader(i);
      group_size[leader_buf[i]]++;
    }
    std::vector<std::vector<int>> result(_n);
    for (int i = 0; i < _n; i++) {
      result[i].reserve(group_size[i]);
    }
    for (int i = 0; i < _n; i++) {
      result[leader_buf[i]].push_back(i);
    }
    result.erase(
        std::remove_if(result.begin(), result.end(),
                       [&](const std::vector<int> &v) { return v.empty(); }),
        result.end());
    return result;
  }

private:
  int _n;
  // root node: -1 * component size
  // otherwise: parent
  std::vector<int> parent_or_size;
};
// =======================================

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  ll N;
  cin >> N;
  dsu d(N);
  vector<vector<ll>> P(N);
  set<ll> y_remain;
  vector<ll> X(N), Y(N);
  vector<ll> y_to_index(N + 1);
  for (ll i = 0; i < N; ++i) {
    cin >> X.at(i) >> Y.at(i);
    P.at(i) = {X.at(i), Y.at(i), i};
    y_to_index.at(Y.at(i)) = i;
    y_remain.insert(Y.at(i));
  }
  sort(P.begin(), P.end());
  for (ll i = 0; i < N; ++i) {
    vector<ll> now = P.at(i);
    ll x = now.at(0), y = now.at(1), index = now.at(2);
    auto itr1 = y_remain.lower_bound(y);
    if (itr1 == y_remain.end() || *itr1 != y) {
      continue;
    }
    y_remain.erase(y);
    auto itr = y_remain.lower_bound(y);
    while (itr != y_remain.end()) {
      ll temp_y = *itr;
      if (d.same(y_to_index.at(y), y_to_index.at(temp_y))) {
        break;
      }
      d.merge(y_to_index.at(y), y_to_index.at(temp_y));
      // cout << y_to_index.at(y) << ' ' << y_to_index.at(temp_y) << "\\n";
      itr++;
    }
  }
  for (ll i = 0; i < N; ++i) {
    cout << d.size(i) << "\\n";
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
// #include <atcoder/all>
// using namespace atcoder;
#define INF 2000000000000000000
#define ll long long
#define pll pair<ll, ll>
using namespace std;

// ===========Disjoint Set Union================
struct dsu {
public:
  dsu() : _n(0) {}
  dsu(int n) : _n(n), parent_or_size(n, -1) {}

  int merge(int a, int b) {
    assert(0 <= a && a < _n);
    assert(0 <= b && b < _n);
    int x = leader(a), y = leader(b);
    if (x == y)
      return x;
    if (-parent_or_size[x] < -parent_or_size[y])
      std::swap(x, y);
    parent_or_size[x] += parent_or_size[y];
    parent_or_size[y] = x;
    return x;
  }

  bool same(int a, int b) {
    assert(0 <= a && a < _n);
    assert(0 <= b && b < _n);
    return leader(a) == leader(b);
  }

  int leader(int a) {
    assert(0 <= a && a < _n);
    if (parent_or_size[a] < 0)
      return a;
    return parent_or_size[a] = leader(parent_or_size[a]);
  }

  int size(int a) {
    assert(0 <= a && a < _n);
    return -parent_or_size[leader(a)];
  }

  std::vector<std::vector<int>> groups() {
    std::vector<int> leader_buf(_n), group_size(_n);
    for (int i = 0; i < _n; i++) {
      leader_buf[i] = leader(i);
      group_size[leader_buf[i]]++;
    }
    std::vector<std::vector<int>> result(_n);
    for (int i = 0; i < _n; i++) {
      result[i].reserve(group_size[i]);
    }
    for (int i = 0; i < _n; i++) {
      result[leader_buf[i]].push_back(i);
    }
    result.erase(
        std::remove_if(result.begin(), result.end(),
                       [&](const std::vector<int> &v) { return v.empty(); }),
        result.end());
    return result;
  }

private:
  int _n;
  // root node: -1 * component size
  // otherwise: parent
  std::vector<int> parent_or_size;
};
// =======================================

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  ll N;
  cin >> N;
  dsu d(N);
  vector<vector<ll>> P(N);
  set<ll> y_remain;
  vector<ll> X(N), Y(N);
  vector<ll> y_to_index(N + 1);
  for (ll i = 0; i < N; ++i) {
    cin >> X.at(i) >> Y.at(i);
    P.at(i) = {X.at(i), Y.at(i), i};
    y_to_index.at(Y.at(i)) = i;
    y_remain.insert(Y.at(i));
  }
  sort(P.begin(), P.end());
  for (ll i = 0; i < N; ++i) {
    vector<ll> now = P.at(i);
    ll x = now.at(0), y = now.at(1), index = now.at(2);
    auto itr1 = y_remain.lower_bound(y);
    if (itr1 == y_remain.end() || *itr1 != y) {
      continue;
    }
    y_remain.erase(y);
    auto itr = y_remain.lower_bound(y);
    while (itr != y_remain.end()) {
      ll temp_y = *itr;
      d.merge(y_to_index.at(y), y_to_index.at(temp_y));
      itr++;
    }
  }
  for (ll i = 0; i < N; ++i) {
    cout << d.size(i) << "\\n";
  }
}