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

#include <bits/stdc++.h>
using namespace std;
const int INF = 1 << 30;

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &a) {
  return is >> a.first >> a.second;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, pair<T1, T2> &a) {
  return os << a.first << " " << a.second;
}
template <typename T> istream &operator>>(istream &is, vector<T> &vc) {
  for (int i = 0; i < vc.size(); i++)
    is >> vc[i];
  return is;
}
template <typename T> ostream &operator<<(ostream &os, vector<T> &vc) {
  for (int i = 0; i < vc.size(); i++)
    os << vc[i] << endl;
  return os;
}

#define ForEach(it, c)                                                         \\
  for (__typeof(c).begin() it = (c).begin(); it != (c).end(); it++)
#define ALL(v) (v).begin(), (v).end()
#define UNQ(s)                                                                 \\
  {                                                                            \\
    sort(ALL(s));                                                              \\
    (s).erase(unique(ALL(s)), (s).end());                                      \\
  }
#define fr first
#define sc second

typedef pair<int, int> Pi;
typedef pair<bool, Pi> Pii;

typedef long long int64;

int UF[40000];
int BF[40000];
int chg[400000];
int siz;
int FIND(int a) {
  if (UF[a] < 0)
    return a;
  return FIND(UF[a]);
}
void UNION(int a, int b) {
  a = FIND(a);
  b = FIND(b);
  if (a == b)
    return;
  if (UF[a] > UF[b])
    swap(a, b);
  chg[siz++] = a;
  chg[siz++] = b;
  UF[a] += UF[b];
  UF[b] = a;
}

const int SQR = 200;
set<Pi> T; // 変化があった辺
set<Pi> S; // 存在する辺

int main() {
  int n, Q;
  int x[40000], y[40000], z[40000];

  scanf("%d %d", &n, &Q);
  for (int i = 0; i < Q; i++) {
    scanf("%d %d %d", &x[i], &y[i], &z[i]);
  }
  for (int i = 0; i < Q; i++) {
    if (i % SQR == 0) {
      T.clear();
      for (int j = 0; j < n; ++j)
        UF[j] = BF[j] = -1;
      for (int j = 0; j < SQR && i + j < Q; j++) {
        if (x[i + j] != 3) {
          T.insert(make_pair(y[i + j], z[i + j]));
          if (x[i + j] == 1)
            S.insert(make_pair(y[i + j], z[i + j]));
          else if (x[i + j] == 2)
            S.erase(make_pair(y[i + j], z[i + j]));
        }
      }
      siz = 0;
      ForEach(it, S) {
        if (T.find(*it) == T.end()) { // 以前に追加されていて変化がない辺を結合
          UNION(it->first, it->second);
        }
      }
      for (int j = 0; j < n; ++j)
        BF[j] = UF[j];
    }
    if (x[i] == 3) {
      set<Pi> u;
      int N = i / SQR * SQR;
      for (int k = 0; k < i % SQR; k++) { // それ以前
        if (x[k + N] == 1)
          u.insert(make_pair(y[N + k], z[N + k]));
        else if (x[k + N] == 2)
          u.erase(make_pair(y[N + k], z[N + k]));
      }
      for (int k = min(SQR - 1, Q - N - 1); k > i % SQR; k--) { // それ以後
        if (x[k + N] == 1)
          u.erase(make_pair(y[N + k], z[N + k]));
        else if (x[k + N] == 2)
          u.insert(make_pair(y[N + k], z[N + k]));
      }
      siz = 0;
      ForEach(it, u) { UNION(it->first, it->second); }
      puts(FIND(y[i]) == FIND(z[i]) ? "YES" : "NO");
      for (int j = 0; j < siz; j++)
        UF[chg[j]] = BF[chg[j]];
    }
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
const int INF = 1 << 30;

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &a) {
  return is >> a.first >> a.second;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, pair<T1, T2> &a) {
  return os << a.first << " " << a.second;
}
template <typename T> istream &operator>>(istream &is, vector<T> &vc) {
  for (int i = 0; i < vc.size(); i++)
    is >> vc[i];
  return is;
}
template <typename T> ostream &operator<<(ostream &os, vector<T> &vc) {
  for (int i = 0; i < vc.size(); i++)
    os << vc[i] << endl;
  return os;
}

#define ForEach(it, c)                                                         \\
  for (__typeof(c).begin() it = (c).begin(); it != (c).end(); it++)
#define ALL(v) (v).begin(), (v).end()
#define UNQ(s)                                                                 \\
  {                                                                            \\
    sort(ALL(s));                                                              \\
    (s).erase(unique(ALL(s)), (s).end());                                      \\
  }
#define fr first
#define sc second

typedef pair<int, int> Pi;
typedef pair<int, Pi> Pii;

typedef long long int64;

class UnionFind {
private:
  vector<int> rank, p;
  stack<Pii> log;

  void link(int x, int y) {
    if (rank[x] > rank[y]) {
      p[y] = x;
    } else {
      p[x] = y;
      if (rank[x] == rank[y])
        rank[y]++;
    }
  }

public:
  UnionFind(){};
  void init(int size) {
    rank.assign(size, 0);
    p.assign(size, 0);
    for (int i = 0; i < size; i++)
      p[i] = i, rank[i] = 0;
  }

  void Union(int x, int y, bool flag = false) {
    if (Find(x) != Find(y)) {
      if (flag)
        log.push(Pii(rank[x] == rank[y], Pi(x, y)));
      link(Find(x), Find(y));
    }
  }
  int Find(int x) { return (x != p[x] ? p[x] = Find(p[x]) : p[x]); }
  void BackUp() {
    while (!log.empty()) {
      Pii top = log.top();
      log.pop();
      p[top.second.first] = top.second.first;
      p[top.second.second] = top.second.second;
      if (top.first)
        rank[top.second.second]--;
    }
  }
};
const int SQR = 150;
set<Pi> T; // 変化があった辺
set<Pi> S; // 存在する辺

int main() {
  int n, Q;
  int x[40000], y[40000], z[40000];

  scanf("%d %d", &n, &Q);
  UnionFind uf;
  for (int i = 0; i < Q; i++) {
    scanf("%d %d %d", &x[i], &y[i], &z[i]);
  }
  for (int i = 0; i * SQR < Q; i++) {
    int begin = i * SQR, end = min((i + 1) * SQR, Q);
    T.clear();
    uf.init(n);
    for (int j = begin; j < end; j++) {
      if (x[j] != 3) {
        T.insert(make_pair(y[j], z[j]));
        if (x[j] == 1)
          S.insert(make_pair(y[j], z[j]));
        else if (x[j] == 2)
          S.erase(make_pair(y[j], z[j]));
      }
    }
    ForEach(it, S) {
      if (T.find(*it) == T.end()) { // 以前に追加されていて変化がない辺を結合
        uf.Union(it->first, it->second);
      }
    }

    for (int j = begin; j < end; j++) {
      if (x[j] == 3) {
        set<Pi> u;
        for (int k = begin; k < j % SQR; k++) { // それ以前
          if (x[k] == 1)
            u.insert(make_pair(y[k], z[k]));
          else if (x[k] == 2)
            u.erase(make_pair(y[k], z[k]));
        }
        for (int k = end - 1; k > j % SQR; k--) { // それ以後
          if (x[k] == 1)
            u.erase(make_pair(y[k], z[k]));
          else if (x[k] == 2)
            u.insert(make_pair(y[k], z[k]));
        }
        ForEach(it, u) { uf.Union(it->first, it->second, true); }
        puts(uf.Find(y[j]) == uf.Find(z[j]) ? "YES" : "NO");
        uf.BackUp();
      }
    }
  }
}