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

#include <bits/stdc++.h>
using namespace std;

#define for_(i, a, b) for (int i = (a); i < (b); ++i)
#define allof(a) (a).begin(), (a).end()
typedef pair<int, int> pii;

struct Edge : pii {
  Edge(int u, int v) : pii(u, v) {}
  int u() { return first; }
  int v() { return second; }
};

class UnionFind {
private:
  vector<int> data;
  int num_of_g;

public:
  UnionFind(int n) : data(n, -1), num_of_g(n) {}

  bool unionSet(int x, int y) {
    x = root(x);
    y = root(y);

    if (x == y)
      return 0;

    if (data[y] < data[x])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    --num_of_g;

    return 1;
  }

  bool sameSet(int x, int y) { return root(x) == root(y); }
  int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
  int size(int x) { return -data[root(x)]; }
  int getNumOfG() { return num_of_g; }
};

#define BUCKET_SIZE 25
int n, k, id[40010], t[40010], u[40010], v[40010];
vector<multiset<int>> eset;
vector<int> vis;

bool dfs(int v, int tar, int col) {
  if (v == tar)
    return true;
  if (vis[v] == col)
    return false;

  vis[v] = col;

  for (int u : eset[v])
    if (dfs(u, tar, col))
      return true;

  return false;
}

void solve() {
  set<Edge> edges;

  vis.assign(n, -1);
  int cur_col = 0;

  for (int i = 0; i < k; i += BUCKET_SIZE) {
    // vertex setting
    vector<int> vset;

    for (int rep = 0; i + rep < k && rep < BUCKET_SIZE; ++rep) {
      int q = i + rep;
      vset.push_back(u[q]);
      vset.push_back(v[q]);
    }

    sort(allof(vset));
    vset.erase(unique(allof(vset)), vset.end());

    // graph construction
    int V = vset.size();

    for_(i, 0, n) id[i] = -1;
    for_(i, 0, V) id[vset[i]] = i;

    eset.assign(V, multiset<int>());

    UnionFind uf(n);

    for (Edge e : edges) {
      if (id[e.u()] != -1 && id[e.v()] != -1) {
        eset[id[e.u()]].insert(id[e.v()]);
        eset[id[e.v()]].insert(id[e.u()]);
      } else {
        uf.unionSet(e.u(), e.v());
      }
    }

    for_(i, 0, V) for_(j, i + 1, V) {
      int u = vset[i], v = vset[j];

      if (uf.sameSet(u, v)) {
        eset[i].insert(j);
        eset[j].insert(i);
      }
    }

    // query
    for (int rep = 0; i + rep < k && rep < BUCKET_SIZE; ++rep) {
      int q = i + rep;

      if (t[q] == 1) {
        eset[id[u[q]]].insert(id[v[q]]);
        eset[id[v[q]]].insert(id[u[q]]);
        edges.insert(Edge(u[q], v[q]));
      }

      if (t[q] == 2) {
        eset[id[u[q]]].erase(id[v[q]]);
        eset[id[v[q]]].erase(id[u[q]]);
        edges.erase(Edge(u[q], v[q]));
      }

      if (t[q] == 3) {
        cout << (dfs(id[u[q]], id[v[q]], cur_col) ? "YES" : "NO") << endl;
        ++cur_col;
      }
    }
  }
}

int main() {
  cin >> n >> k;
  for_(i, 0, k) cin >> t[i] >> u[i] >> v[i];
  solve();
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define for_(i, a, b) for (int i = (a); i < (b); ++i)
#define allof(a) (a).begin(), (a).end()
typedef pair<int, int> pii;

struct Edge : pii {
  Edge(int u, int v) : pii(u, v) {}
  int u() { return first; }
  int v() { return second; }
};

class UnionFind {
private:
  vector<int> data;
  int num_of_g;

public:
  UnionFind(int n) : data(n, -1), num_of_g(n) {}

  bool unionSet(int x, int y) {
    x = root(x);
    y = root(y);

    if (x == y)
      return 0;

    if (data[y] < data[x])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    --num_of_g;

    return 1;
  }

  bool sameSet(int x, int y) { return root(x) == root(y); }
  int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
  int size(int x) { return -data[root(x)]; }
  int getNumOfG() { return num_of_g; }
};

#define BUCKET_SIZE 50
int n, k, id[40010], t[40010], u[40010], v[40010];
vector<multiset<int>> eset;
vector<int> vis;

bool dfs(int v, int tar, int col) {
  if (v == tar)
    return true;
  if (vis[v] == col)
    return false;

  vis[v] = col;

  for (int u : eset[v])
    if (dfs(u, tar, col))
      return true;

  return false;
}

void solve() {
  set<Edge> edges;

  vis.assign(n, -1);
  int cur_col = 0;

  for (int i = 0; i < k; i += BUCKET_SIZE) {
    // vertex setting
    vector<int> vset;

    for (int rep = 0; i + rep < k && rep < BUCKET_SIZE; ++rep) {
      int q = i + rep;
      vset.push_back(u[q]);
      vset.push_back(v[q]);
    }

    sort(allof(vset));
    vset.erase(unique(allof(vset)), vset.end());

    // graph construction
    int V = vset.size();

    for_(i, 0, n) id[i] = -1;
    for_(i, 0, V) id[vset[i]] = i;

    eset.assign(V, multiset<int>());

    UnionFind uf(n);

    for (Edge e : edges) {
      if (id[e.u()] != -1 && id[e.v()] != -1) {
        eset[id[e.u()]].insert(id[e.v()]);
        eset[id[e.v()]].insert(id[e.u()]);
      } else {
        uf.unionSet(e.u(), e.v());
      }
    }

    for_(i, 0, V) for_(j, i + 1, V) {
      int u = vset[i], v = vset[j];

      if (uf.sameSet(u, v)) {
        eset[i].insert(j);
        eset[j].insert(i);
      }
    }

    // query
    for (int rep = 0; i + rep < k && rep < BUCKET_SIZE; ++rep) {
      int q = i + rep;

      if (t[q] == 1) {
        eset[id[u[q]]].insert(id[v[q]]);
        eset[id[v[q]]].insert(id[u[q]]);
        edges.insert(Edge(u[q], v[q]));
      }

      if (t[q] == 2) {
        eset[id[u[q]]].erase(id[v[q]]);
        eset[id[v[q]]].erase(id[u[q]]);
        edges.erase(Edge(u[q], v[q]));
      }

      if (t[q] == 3) {
        cout << (dfs(id[u[q]], id[v[q]], cur_col) ? "YES" : "NO") << endl;
        ++cur_col;
      }
    }
  }
}

int main() {
  cin >> n >> k;
  for_(i, 0, k) cin >> t[i] >> u[i] >> v[i];
  solve();
}