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

#include <iostream>
#include <map>
#include <set>
#include <vector>

using namespace std;
using ll = long long;

struct DisjointSets {
  DisjointSets(const int v) {
    m_parent.resize(v);
    m_rank.resize(v);
    for (int i = 0; i < v; i++) {
      m_parent[i] = i;
      m_rank[i] = 0;
    }
  }

  bool same(const int a, const int b) { return find(a) == find(b); }

  int find(const int a) const {
    if (m_parent[a] == a) {
      return a;
    } else {
      return find(m_parent[a]);
    }
  }

  void unite(const int a_, const int b_) {
    const int a = find(a_);
    const int b = find(b_);
    if (a == b) {
      return;
    }
    if (m_rank[a] > m_rank[b]) {
      m_parent[b] = a;
    } else {
      m_parent[a] = b;
    }
    if (m_rank[a] == m_rank[b]) {
      m_rank[b]++;
    }
  }

  vector<int> m_parent;
  vector<int> m_rank;
};

struct Query {
  int q;
  int u;
  int v;
};

int main() {
  int N, K;
  cin >> N >> K;
  const int SQRT = 200;
  const int BNUM = (K + SQRT - 1) / SQRT;
  using P = pair<int, int>;
  vector<vector<Query>> queries(BNUM);
  for (int Q = 0; Q < K; Q++) {
    int q, u, v;
    cin >> q >> u >> v;
    if (u > v) {
      swap(u, v);
    }
    queries[Q / SQRT].push_back(Query{q, u, v});
  }
  for (int i = 0; i < BNUM; i++) {
    set<P> edge;
    for (int j = 0; j < i; j++) {
      for (const auto &q : queries[j]) {
        if (q.q == 1) {
          edge.insert(make_pair(q.u, q.v));
        } else if (q.q == 2) {
          edge.erase(make_pair(q.u, q.v));
        }
      }
    }
    set<P> in;
    set<int> node;
    for (const auto &q : queries[i]) {
      if (q.q == 2) {
        if (edge.find(make_pair(q.u, q.v)) != edge.end()) {
          edge.erase(make_pair(q.u, q.v));
          in.insert(make_pair(q.u, q.v));
        }
        node.insert(q.u);
        node.insert(q.v);
      } else if (q.q == 1) {
        node.insert(q.u);
        node.insert(q.v);
      }
    }

    DisjointSets uf(N);
    for (const auto &e : edge) {
      uf.unite(e.first, e.second);
    }
    auto uf_ = uf;

    for (const auto &q : queries[i]) {
      if (q.q == 1) {
        in.insert(make_pair(q.u, q.v));
      } else if (q.q == 2) {
        in.erase(make_pair(q.u, q.v));
      } else {
        for (const auto &e : in) {
          node.insert(e.first);
          node.insert(e.second);
          node.insert(uf.find(e.first));
          node.insert(uf.find(e.second));
          uf.unite(e.first, e.second);
        }
        cout << (uf.same(q.u, q.v) ? "YES" : "NO") << endl;
        for (const int n : node) {
          uf.m_parent[n] = uf_.m_parent[n];
        }
      }
    }
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <iostream>
#include <set>
#include <vector>

using namespace std;
using ll = long long;

class DisjointSets {
public:
  DisjointSets(const int v) {
    m_parent.resize(v);
    m_rank.resize(v);
    for (int i = 0; i < v; i++) {
      m_parent[i] = i;
      m_rank[i] = 0;
    }
  }

  bool same(const int a, const int b) { return find(a) == find(b); }

  int find(const int a) {
    if (m_parent[a] == a) {
      return a;
    } else {
      return m_parent[a] = find(m_parent[a]);
    }
  }

  void unite(const int a_, const int b_) {
    const int a = find(a_);
    const int b = find(b_);
    if (a == b) {
      return;
    }
    if (m_rank[a] > m_rank[b]) {
      m_parent[b] = a;
    } else {
      m_parent[a] = b;
    }
    if (m_rank[a] == m_rank[b]) {
      m_rank[b]++;
    }
  }

private:
  vector<int> m_parent;
  vector<int> m_rank;
};

struct Query {
  int q;
  int u;
  int v;
};

int main() {
  int N, K;
  cin >> N >> K;
  const int SQRT = 200;
  const int BNUM = (K + SQRT - 1) / SQRT;
  using P = pair<int, int>;
  vector<vector<Query>> queries(BNUM);
  for (int Q = 0; Q < K; Q++) {
    int q, u, v;
    cin >> q >> u >> v;
    if (u > v) {
      swap(u, v);
    }
    queries[Q / SQRT].push_back(Query{q, u, v});
  }
  for (int i = 0; i < BNUM; i++) {
    set<P> edge;
    for (int j = 0; j < i; j++) {
      for (const auto &q : queries[j]) {
        if (q.q == 1) {
          edge.insert(make_pair(q.u, q.v));
        } else if (q.q == 2) {
          edge.erase(make_pair(q.u, q.v));
        }
      }
    }
    set<P> in;
    for (const auto &q : queries[i]) {
      if (q.q == 2) {
        if (edge.find(make_pair(q.u, q.v)) != edge.end()) {
          edge.erase(make_pair(q.u, q.v));
          in.insert(make_pair(q.u, q.v));
        }
      }
    }
    DisjointSets uf(N);
    for (const auto &e : edge) {
      uf.unite(e.first, e.second);
    }

    for (const auto &q : queries[i]) {
      if (q.q == 1) {
        in.insert(make_pair(q.u, q.v));
      } else if (q.q == 2) {
        in.erase(make_pair(q.u, q.v));
      } else {
        auto UF = uf;
        for (const auto &e : in) {
          UF.unite(e.first, e.second);
        }
        cout << (UF.same(q.u, q.v) ? "YES" : "NO") << endl;
      }
    }
  }

  return 0;
}