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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <map>
#include <set>
#include <vector>
using namespace std;

#define INF 1e9

// UnionFind??¨ ??????????????????
class UnionFindTree {
public:
  UnionFindTree();
  UnionFindTree(int size);
  void init(int size);
  void unite(int a, int b);
  bool same(int a, int b);
  bool isConnected(void);

private:
  int m_size, unite_count;
  vector<int> parent;
  vector<int> rank;
  int find(int a);
};

// ?????????????????????
UnionFindTree::UnionFindTree(){};
UnionFindTree::UnionFindTree(int size) { init(size); };

// UF??¨????????????
void UnionFindTree::init(int size) {
  m_size = size;
  unite_count = 0;
  parent.clear();
  rank.clear();
  for (int i = 0; i < size; i++) {
    parent.push_back(i);
    rank.push_back(0);
  }
  return;
}

// UF??¨??§???????????£?????????
void UnionFindTree::unite(int a, int b) {
  int x = find(a);
  int y = find(b);
  // ?????§???????????£???????????????????????????
  if (x == y) {
    return;
  }

  unite_count++;
  if (rank[x] < rank[y]) {
    parent[x] = y;
  } else {
    parent[y] = x;
    if (rank[x] == rank[y]) {
      rank[x]++;
    }
  }
  return;
}

// UF??¨??§???????????????????????????????????????
bool UnionFindTree::same(int a, int b) { return find(a) == find(b); }

// UF??¨??§?????????????????????????????¨???????????????????????????
bool UnionFindTree::isConnected(void) { return unite_count == m_size - 1; }

// UF??¨?????????????????????
int UnionFindTree::find(int a) {
  if (parent[a] != a) {
    parent[a] = find(parent[a]);
  }
  return parent[a];
}

bool dfs(int now, int goal, const vector<vector<int>> &adj,
         vector<bool> &used) {
  used[now] = true;
  if (now == goal) {
    return true;
  }
  for (size_t i = 0; i < adj[now].size(); i++)
    if (!used[adj[now][i]] && dfs(adj[now][i], goal, adj, used))
      return true;
  return false;
}

int main(void) {
  int n, k;
  scanf("%d%d", &n, &k);

  int root_k = max((int)sqrt(k), 1);
  int rk_size = ((k - 1) / root_k) + 1;
  vector<vector<pair<int, pair<int, int>>>> query(rk_size);
  vector<set<int>> attention_vertex(rk_size);
  vector<pair<pair<int, int>, pair<int, int>>> edge;
  map<pair<int, int>, int> nowEdge;

  // O(klogk)
  for (int i = 0; i < k; i++) {
    // ??\\??????????????????
    int q, u, v;
    scanf("%d%d%d", &q, &u, &v);

    // rk_size???????????????????????????????????????????????????????????¨?????¨?????????????????§????????????????????????????????????
    // O(log???k)
    int at = i / root_k;
    query[at].push_back(make_pair(q, make_pair(u, v)));
    attention_vertex[at].insert(u);
    attention_vertex[at].insert(v);

    // ??????????????¨?¶???????????????????
    // nowEdge??????????????§?????¨??????????????????
    // O(logk)
    if (q == 1) {
      nowEdge[make_pair(u, v)] = i;
    } else if (q == 2) {
      // ???????¶??????????????????????????????????????????¨?¶???±??????????????¢??????vector?????????
      int s = nowEdge[make_pair(u, v)];
      nowEdge.erase(make_pair(u, v));
      edge.push_back(make_pair(make_pair(u, v), make_pair(s, i)));
    }
  }
  // ???????????§?¶???????????????£??????????????????????????????
  for (map<pair<int, int>, int>::iterator it = nowEdge.begin();
       it != nowEdge.end(); it++)
    edge.push_back(make_pair((*it).first, make_pair((*it).second, INF)));

  // ??????????????????????????¨??????????????????
  // O(k???ka(n) + n???k)
  for (int i = 0; i < (int)query.size(); i++) {
    // ??????????????¨?????£??\\??????????????£?¨?
    // O(n)
    vector<vector<int>> adj(n);

    // ????????????????????£??¨??£?????§???????????????UnionFindTree??§????????????
    // O(n + ka(n))
    UnionFindTree uf(n);
    for (size_t j = 0; j < edge.size(); j++) {
      if (edge[j].second.first <= (root_k * i) &&
          (root_k * (i + 1)) <= edge[j].second.second) {
        uf.unite(edge[j].first.first, edge[j].first.second);
      } else if (edge[j].second.first < (root_k * i) &&
                 (root_k * i) <= edge[j].second.second &&
                 edge[j].second.second < (root_k * (i + 1))) {
        adj[edge[j].first.first].push_back(edge[j].first.second);
        adj[edge[j].first.second].push_back(edge[j].first.first);
      }
    }

    // UnionFindTree????????±??????????????????????????¨?????£??\\??????????????´??°
    // O(k)
    for (set<int>::iterator it = attention_vertex[i].begin();
         it != attention_vertex[i].end(); it++)
      for (set<int>::iterator te = attention_vertex[i].begin();
           te != attention_vertex[i].end(); te++)
        if (((*it) != (*te)) && uf.same(*it, *te))
          adj[*it].push_back(*te);

    // ?????¨?????????
    // O(k)
    for (size_t j = 0; j < query[i].size(); j++) {
      int q = query[i][j].first;
      int u = query[i][j].second.first;
      int v = query[i][j].second.second;

      switch (q) {
      case 1:
        adj[u].push_back(v);
        adj[v].push_back(u);
        break;
      case 2:
        // O(???k)
        for (int k = adj[u].size() - 1; k >= 0; k--) {
          if (adj[u][k] == v) {
            adj[u].erase(adj[u].begin() + k);
            break;
          }
        }
        for (int k = adj[v].size() - 1; k >= 0; k--) {
          if (adj[v][k] == u) {
            adj[v].erase(adj[v].begin() + k);
            break;
          }
        }
        break;
      case 3:
        // O(???k)
        vector<bool> used(adj.size(), false);
        if (dfs(u, v, adj, used)) {
          printf("YES\\n");
        } else {
          printf("NO\\n");
        }
        break;
      }
    }
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <set>
#include <vector>
using namespace std;

#define INF 1e9

// UnionFind??¨ ??????????????????
class UnionFindTree {
public:
  UnionFindTree();
  UnionFindTree(int size);
  void init(int size);
  void unite(int a, int b);
  bool same(int a, int b);
  bool isConnected(void);

private:
  int m_size, unite_count;
  vector<int> parent;
  vector<int> rank;
  int find(int a);
};

// ?????????????????????
UnionFindTree::UnionFindTree(){};
UnionFindTree::UnionFindTree(int size) { init(size); };

// UF??¨????????????
void UnionFindTree::init(int size) {
  m_size = size;
  unite_count = 0;
  parent.clear();
  rank.clear();
  for (int i = 0; i < size; i++) {
    parent.push_back(i);
    rank.push_back(0);
  }
  return;
}

// UF??¨??§???????????£?????????
void UnionFindTree::unite(int a, int b) {
  int x = find(a);
  int y = find(b);
  // ?????§???????????£???????????????????????????
  if (x == y) {
    return;
  }

  unite_count++;
  if (rank[x] < rank[y]) {
    parent[x] = y;
  } else {
    parent[y] = x;
    if (rank[x] == rank[y]) {
      rank[x]++;
    }
  }
  return;
}

// UF??¨??§???????????????????????????????????????
bool UnionFindTree::same(int a, int b) { return find(a) == find(b); }

// UF??¨??§?????????????????????????????¨???????????????????????????
bool UnionFindTree::isConnected(void) { return unite_count == m_size - 1; }

// UF??¨?????????????????????
int UnionFindTree::find(int a) {
  if (parent[a] != a) {
    parent[a] = find(parent[a]);
  }
  return parent[a];
}

bool dfs(int now, int goal, const vector<vector<int>> &adj,
         vector<bool> &used) {
  used[now] = true;
  if (now == goal) {
    return true;
  }
  for (size_t i = 0; i < adj[now].size(); i++)
    if (!used[adj[now][i]] && dfs(adj[now][i], goal, adj, used))
      return true;
  return false;
}

int main(void) {
  int n, k;
  cin >> n >> k;

  int root_k = max((int)sqrt(k), 1);
  int rk_size = ((k - 1) / root_k) + 1;
  vector<vector<pair<int, pair<int, int>>>> query(rk_size);
  vector<set<int>> attention_vertex(rk_size);
  vector<pair<pair<int, int>, pair<int, int>>> edge;
  map<pair<int, int>, int> nowEdge;

  // O(klogk)
  for (int i = 0; i < k; i++) {
    // ??\\??????????????????
    int q, u, v;
    cin >> q >> u >> v;

    // rk_size???????????????????????????????????????????????????????????¨?????¨?????????????????§????????????????????????????????????
    // O(log???k)
    int at = i / root_k;
    query[at].push_back(make_pair(q, make_pair(u, v)));
    attention_vertex[at].insert(u);
    attention_vertex[at].insert(v);

    // ??????????????¨?¶???????????????????
    // nowEdge??????????????§?????¨??????????????????
    // O(logk)
    if (q == 1) {
      nowEdge[make_pair(u, v)] = i;
    } else if (q == 2) {
      // ???????¶??????????????????????????????????????????¨?¶???±??????????????¢??????vector?????????
      int s = nowEdge[make_pair(u, v)];
      nowEdge.erase(make_pair(u, v));
      edge.push_back(make_pair(make_pair(u, v), make_pair(s, i)));
    }
  }
  // ???????????§?¶???????????????£??????????????????????????????
  for (map<pair<int, int>, int>::iterator it = nowEdge.begin();
       it != nowEdge.end(); it++)
    edge.push_back(make_pair((*it).first, make_pair((*it).second, INF)));

  // ??????????????????????????¨??????????????????
  // O(k???ka(n) + n???k)
  for (int i = 0; i < (int)query.size(); i++) {
    // ??????????????¨?????£??\\??????????????£?¨?
    // O(n)
    vector<vector<int>> adj(n);

    // ????????????????????£??¨??£?????§???????????????UnionFindTree??§????????????
    // O(n + ka(n))
    UnionFindTree uf(n);
    for (size_t j = 0; j < edge.size(); j++) {
      if (edge[j].second.first <= (root_k * i) &&
          (root_k * (i + 1)) <= edge[j].second.second) {
        uf.unite(edge[j].first.first, edge[j].first.second);
      } else if (edge[j].second.first < (root_k * i) &&
                 (root_k * i) <= edge[j].second.second &&
                 edge[j].second.second < (root_k * (i + 1))) {
        adj[edge[j].first.first].push_back(edge[j].first.second);
        adj[edge[j].first.second].push_back(edge[j].first.first);
      }
    }

    // UnionFindTree????????±??????????????????????????¨?????£??\\??????????????´??°
    // O(k)
    for (set<int>::iterator it = attention_vertex[i].begin();
         it != attention_vertex[i].end(); it++)
      for (set<int>::iterator te = attention_vertex[i].begin();
           te != attention_vertex[i].end(); te++)
        if (((*it) != (*te)) && uf.same(*it, *te))
          adj[*it].push_back(*te);

    // ?????¨?????????
    // O(k)
    for (size_t j = 0; j < query[i].size(); j++) {
      int q = query[i][j].first;
      int u = query[i][j].second.first;
      int v = query[i][j].second.second;

      switch (q) {
      case 1:
        adj[u].push_back(v);
        adj[v].push_back(u);
        break;
      case 2:
        // O(???k)
        for (int k = adj[u].size() - 1; k >= 0; k--) {
          if (adj[u][k] == v) {
            adj[u].erase(adj[u].begin() + k);
            break;
          }
        }
        for (int k = adj[v].size() - 1; k >= 0; k--) {
          if (adj[v][k] == u) {
            adj[v].erase(adj[v].begin() + k);
            break;
          }
        }
        break;
      case 3:
        // O(???k)
        vector<bool> used(adj.size(), false);
        if (dfs(u, v, adj, used)) {
          cout << "YES" << endl;
        } else {
          cout << "NO" << endl;
        }
        break;
      }
    }
  }

  return 0;
}