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

#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #pragma GCC optimize ("trapv")
using namespace std;

#define pii pair<int, int>
#define pll pair<long long, long long>
#define fi first
#define se second
#define mp make_pair
#define pb push_back

#define DEBUG 1
#define cerr                                                                   \\
  if (DEBUG)                                                                   \\
  cerr
#define test cerr << "hi\\n";

// mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
// #define rand(l, r) uniform_int_distribution<long long>(l, r)(rng)
// using namespace __gnu_pbds;
// typedef tree<int, null_type, less<int>, rb_tree_tag,
// tree_order_statistics_node_update> OST;

#define INF 0x3f3f3f3f3f3f3f3f
#define MAXN 200005
// #define MOD 998244353LL
#define MOD 1000000007LL

int n, m;
vector<pair<pii, int>> games;
int p[MAXN], col[MAXN];
vector<int> v[MAXN];

void join(int x, int y, int st) {
  x = p[x], y = p[y];
  if (x == y)
    return;
  if (v[x].size() < v[y].size())
    swap(x, y);

  // v[x] > v[y]

  int col_shift = col[x] + st + 3 - col[y];

  int src = p[y];
  for (int i = 0; i < v[src].size(); i++) {
    int q = v[src][i];
    p[q] = p[x];
    v[p[x]].pb(q);
    (col[q] += col_shift) %= 3;
  }
  v[src].clear();
}

int work(int go) {
  for (int i = 0; i < n; i++) {
    p[i] = i;
    v[i].clear();
    v[i].pb(i);
    col[i] = 0;
  }

  for (int i = 0; i < m; i++) {
    int u, v, st;
    u = games[i].fi.fi;
    v = games[i].fi.se;
    st = games[i].se;

    if (u == go or v == go)
      continue;

    if (p[u] != p[v]) {
      join(u, v, st);
    } else {
      if ((col[u] + st) % 3 != col[v])
        return i;
    }
  }

  return -1;
}

int val[255];
int q;

void solve() {
  cin >> n >> q;
  iota(p + 1, p + 1 + n, 1);
  for (int i = 0; i <= n; i++)
    v[i].pb(i);
  while (q--) {
    int t, u, v;
    cin >> t >> u >> v;
    if (t == 0) {
      join(u, v, 0);
    }
    if (t == 1) {
      cout << (p[u] == p[v]) << '\\n';
    }
  }
}

int main() {
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
  //    freopen("output.txt", "w", stdout);
  auto start_time = clock();
#endif
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int t = 1;
  //    cin >> t;
  for (int no = 1; no <= t; no++) {
    //        cout << "Case #" << no << ": ";
    solve();
  }

#ifdef LOCAL
  cerr << "TIME ELAPSED: " << double(clock() - start_time) / CLOCKS_PER_SEC
       << '\\n';
#endif
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #pragma GCC optimize ("trapv")
using namespace std;

#define pii pair<int, int>
#define pll pair<long long, long long>
#define fi first
#define se second
#define mp make_pair
#define pb push_back

#define DEBUG 1
#define cerr                                                                   \\
  if (DEBUG)                                                                   \\
  cerr
#define test cerr << "hi\\n";

// mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
// #define rand(l, r) uniform_int_distribution<long long>(l, r)(rng)
// using namespace __gnu_pbds;
// typedef tree<int, null_type, less<int>, rb_tree_tag,
// tree_order_statistics_node_update> OST;

#define INF 0x3f3f3f3f3f3f3f3f
#define MAXN 200005
// #define MOD 998244353LL
#define MOD 1000000007LL

int n, m;
vector<pair<pii, int>> games;
int p[MAXN], col[MAXN];
vector<int> v[MAXN];

void join(int x, int y, int st) {
  x = p[x], y = p[y];
  if (v[x].size() < v[y].size())
    swap(x, y);
  while (v[y].size()) {
    int q = v[y].back();
    v[x].pb(q);
    p[q] = x;
    v[y].pop_back();
  }
}

int work(int go) {
  for (int i = 0; i < n; i++) {
    p[i] = i;
    v[i].clear();
    v[i].pb(i);
    col[i] = 0;
  }

  for (int i = 0; i < m; i++) {
    int u, v, st;
    u = games[i].fi.fi;
    v = games[i].fi.se;
    st = games[i].se;

    if (u == go or v == go)
      continue;

    if (p[u] != p[v]) {
      join(u, v, st);
    } else {
      if ((col[u] + st) % 3 != col[v])
        return i;
    }
  }

  return -1;
}

int val[255];
int q;

void solve() {
  cin >> n >> q;
  iota(p + 1, p + 1 + n, 1);
  for (int i = 0; i <= n; i++)
    v[i].pb(i);
  while (q--) {
    int t, u, v;
    cin >> t >> u >> v;
    if (t == 0) {
      join(u, v, 0);
    }
    if (t == 1) {
      cout << (p[u] == p[v]) << '\\n';
    }
  }
}

int main() {
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
  //    freopen("output.txt", "w", stdout);
  auto start_time = clock();
#endif
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int t = 1;
  //    cin >> t;
  for (int no = 1; no <= t; no++) {
    //        cout << "Case #" << no << ": ";
    solve();
  }

#ifdef LOCAL
  cerr << "TIME ELAPSED: " << double(clock() - start_time) / CLOCKS_PER_SEC
       << '\\n';
#endif
}