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

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

const int MAX_N = 40000;
const int MAX_Q = 40000;
const int SQRT = 200;
const int INF = 1 << 29;

struct query {
  int type, x, y;
};

int N, Q, E;
query qs[MAX_Q];

struct edge {
  int from, tar, add, rm;
};

edge es[MAX_Q];
vector<int> graph[MAX_N];
bool found;
bool visited[MAX_N];

struct UnionFind {
  static const int MAX_UF = MAX_N;

  int pars[MAX_UF];
  int ranks[MAX_UF];

  void init(int n) {
    for (int i = 0; i < n; i++) {
      pars[i] = i;
      ranks[i] = 0;
    }
  }

  int get_root(int x) {
    if (pars[x] == x) {
      return x;
    } else {
      return pars[x] = get_root(pars[x]);
    }
  }

  void merge(int x, int y) {
    x = get_root(x);
    y = get_root(y);
    if (x == y) {
      return;
    }

    if (ranks[x] < ranks[y]) {
      pars[x] = y;
    } else if (ranks[y] < ranks[x]) {
      pars[y] = x;
    } else {
      pars[x] = y;
      ranks[y]++;
    }
  }

  bool is_same(int x, int y) { return get_root(x) == get_root(y); }
};

UnionFind uf;

void init() {
  scanf("%d%d", &N, &Q);
  for (int i = 0; i < Q; i++) {
    scanf("%d%d%d", &qs[i].type, &qs[i].x, &qs[i].y);
  }
}

void prepare() {
  map<pair<int, int>, int> dict;
  for (int i = 0; i < Q; i++) {
    if (qs[i].type == 1) {
      dict[make_pair(qs[i].x, qs[i].y)] = E;
      es[E++] = (edge){qs[i].x, qs[i].y, i, INF};
    } else if (qs[i].type == 2) {
      es[dict[make_pair(qs[i].x, qs[i].y)]].rm = i;
    }
  }
}

void prepare2(int start) {
  uf.init(N);
  for (int i = 0; i < E; i++) {
    if (es[i].add < start && es[i].rm >= start + SQRT) {
      uf.merge(es[i].from, es[i].tar);
    }
  }

  for (int i = 0; i < N; i++) {
    graph[i].clear();
  }

  for (int i = 0; i < E; i++) {
    if (es[i].add < start && es[i].rm >= start && es[i].rm < start + SQRT) {
      int x = uf.get_root(es[i].from), y = uf.get_root(es[i].tar);
      graph[x].push_back(y);
      graph[y].push_back(x);
    }
  }
}

void dfs(int v, int goal, bool color) {
  if (v == goal) {
    found = true;
  }
  visited[v] = color;

  for (int i = 0; i < (int)graph[v].size(); i++) {
    const int u = graph[v][i];
    if (visited[u] != color) {
      dfs(u, goal, color);
    }
  }
}

bool search(int start, int goal) {
  found = false;
  dfs(start, goal, true);
  dfs(start, goal, false);
  return found;
}

void solve() {
  prepare();
  for (int b = 0; b * SQRT < Q; b++) {
    prepare2(b * SQRT);
    for (int t = b * SQRT; t < (b + 1) * SQRT && t < Q; t++) {
      int x = uf.get_root(qs[t].x), y = uf.get_root(qs[t].y);
      if (qs[t].type == 1) {
        graph[x].push_back(y);
        graph[y].push_back(x);
      } else if (qs[t].type == 2) {
        graph[x].erase(find(graph[x].begin(), graph[x].end(), y));
        graph[y].erase(find(graph[y].begin(), graph[y].end(), x));
      } else if (qs[t].type == 3) {
        if (x == y || search(x, y)) {
          puts("YES");
        } else {
          puts("NO");
        }
      }
    }
  }
}

int main() {
  init();
  solve();

  return 0;
}

Error: TLE

Fixed submission:

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

const int MAX_N = 40000;
const int MAX_Q = 40000;
const int SQRT = 200;
const int INF = 1 << 29;

struct query {
  int type, x, y;
};

int N, Q, E;
query qs[MAX_Q];

struct edge {
  int from, tar, add, rm;
};

edge es[MAX_Q];
vector<pair<int, bool>> graph[MAX_N];
bool visited[MAX_N];
vector<int> ids;

struct UnionFind {
  static const int MAX_UF = MAX_N;

  int pars[MAX_UF];
  int ranks[MAX_UF];

  void init(int n) {
    for (int i = 0; i < n; i++) {
      pars[i] = i;
      ranks[i] = 0;
    }
  }

  int get_root(int x) {
    if (pars[x] == x) {
      return x;
    } else {
      return pars[x] = get_root(pars[x]);
    }
  }

  void merge(int x, int y) {
    x = get_root(x);
    y = get_root(y);
    if (x == y) {
      return;
    }

    if (ranks[x] < ranks[y]) {
      pars[x] = y;
    } else if (ranks[y] < ranks[x]) {
      pars[y] = x;
    } else {
      pars[x] = y;
      ranks[y]++;
    }
  }

  bool is_same(int x, int y) { return get_root(x) == get_root(y); }
};

UnionFind uf;

map<pair<int, int>, int> stamp;

void init() {
  scanf("%d%d", &N, &Q);
  for (int i = 0; i < Q; i++) {
    scanf("%d%d%d", &qs[i].type, &qs[i].x, &qs[i].y);
  }
}

void prepare() {
  map<pair<int, int>, int> dict;
  for (int i = 0; i < Q; i++) {
    if (qs[i].type == 1) {
      dict[make_pair(qs[i].x, qs[i].y)] = E;
      es[E++] = (edge){qs[i].x, qs[i].y, i, INF};
    } else if (qs[i].type == 2) {
      es[dict[make_pair(qs[i].x, qs[i].y)]].rm = i;
    }
  }
}

void prepare2(int start) {
  uf.init(N);
  for (int i = 0; i < E; i++) {
    if (es[i].add < start && es[i].rm >= start + SQRT) {
      uf.merge(es[i].from, es[i].tar);
    }
  }

  ids.clear();
  for (int i = 0; i < N; i++)
    if (uf.get_root(i) == i) {
      ids.push_back(i);
    }

  for (int i = 0; i < N; i++) {
    graph[i].clear();
  }

  stamp.clear();
  for (int i = 0; i < E; i++) {
    if (es[i].add < start && es[i].rm >= start && es[i].rm < start + SQRT) {
      int x = uf.get_root(es[i].from), y = uf.get_root(es[i].tar);
      stamp[make_pair(x, y)] = graph[x].size();
      stamp[make_pair(y, x)] = graph[y].size();
      graph[x].push_back(make_pair(y, true));
      graph[y].push_back(make_pair(x, true));
    }
  }
}

bool dfs(int v, int goal) {
  if (v == goal) {
    return true;
  }
  visited[v] = true;

  for (int i = 0; i < (int)graph[v].size(); i++) {
    const int u = graph[v][i].first;
    if (graph[v][i].second && !visited[u] && dfs(u, goal)) {
      return true;
    }
  }
  return false;
}

bool search(int start, int goal) {
  for (int i = 0; i < (int)ids.size(); i++) {
    visited[ids[i]] = false;
  }

  return dfs(start, goal);
}

void solve() {
  prepare();
  for (int b = 0; b * SQRT < Q; b++) {
    prepare2(b * SQRT);
    for (int t = b * SQRT; t < (b + 1) * SQRT && t < Q; t++) {
      int x = uf.get_root(qs[t].x), y = uf.get_root(qs[t].y);
      if (qs[t].type == 1) {
        stamp[make_pair(x, y)] = graph[x].size();
        stamp[make_pair(y, x)] = graph[y].size();
        graph[x].push_back(make_pair(y, true));
        graph[y].push_back(make_pair(x, true));
      } else if (qs[t].type == 2) {
        graph[x][stamp[make_pair(x, y)]].second = false;
        graph[y][stamp[make_pair(y, x)]].second = false;
      } else if (qs[t].type == 3) {
        if (x == y || search(x, y)) {
          puts("YES");
        } else {
          puts("NO");
        }
      }
    }
  }
}

int main() {
  init();
  solve();

  return 0;
}