Problem p03855 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()
using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
using Graph = vector<vi>; // vi <=> vl
using P = pair<int, int>; // int <=> ll
const int MOD = 1e9 + 7;
const double EPS = 1e-9;

int main() {
  int N, K, L;
  cin >> N >> K >> L;
  Graph road(N), train(N);
  rep(i, K) {
    int p, q;
    cin >> p >> q;
    --p;
    --q;
    road[p].push_back(q);
    road[q].push_back(p);
  }
  rep(i, L) {
    int r, s;
    cin >> r >> s;
    --r;
    --s;
    train[r].push_back(s);
    train[s].push_back(r);
  }

  vector<bool> seen(N, false);
  vi parent_road(N);
  map<int, set<int>> connect_road;

  rep(i, N) {
    if (seen[i])
      continue;
    seen[i] = true;
    queue<int> q;
    connect_road[i].insert(i);
    q.push(i);

    while (!q.empty()) {
      int v = q.front();
      q.pop();
      for (auto nv : road[v]) {
        if (seen[nv])
          continue;
        seen[nv] = true;
        connect_road[i].insert(nv);
        q.push(nv);
      }
    }

    for (auto v : connect_road[i]) {
      parent_road[v] = i;
    }
  }

  seen.assign(N, false);
  vi parent_train(N);
  map<int, set<int>> connect_train;

  rep(i, N) {
    if (seen[i])
      continue;
    seen[i] = true;
    queue<int> q;
    connect_train[i].insert(i);
    q.push(i);

    while (!q.empty()) {
      int v = q.front();
      q.pop();
      for (auto nv : train[v]) {
        if (seen[nv])
          continue;
        seen[nv] = true;
        connect_train[i].insert(nv);
        q.push(nv);
      }
    }

    for (auto v : connect_train[i]) {
      parent_train[v] = i;
    }
  }

  seen.assign(N, false);
  vi ans(N);
  rep(i, N) {
    if (seen[i])
      continue;
    int a = parent_road[i];
    int b = parent_train[i];
    vi connect;

    for (auto v : connect_road[a]) {
      if (parent_train[i] == parent_train[v]) {
        connect.push_back(v);
        seen[v] = true;
      }
    }

    for (auto v : connect) {
      ans[v] = connect.size();
    }
  }

  rep(i, N - 1) cout << ans[i] << ' ';
  cout << ans[N - 1] << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define all(a) (a).begin(), (a).end()
#define rall(a) (a).rbegin(), (a).rend()
using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
using Graph = vector<vi>; // vi <=> vl
using P = pair<int, int>; // int <=> ll
const int MOD = 1e9 + 7;
const double EPS = 1e-9;

int main() {
  int N, K, L;
  cin >> N >> K >> L;
  Graph road(N), train(N);
  rep(i, K) {
    int p, q;
    cin >> p >> q;
    --p;
    --q;
    road[p].push_back(q);
    road[q].push_back(p);
  }
  rep(i, L) {
    int r, s;
    cin >> r >> s;
    --r;
    --s;
    train[r].push_back(s);
    train[s].push_back(r);
  }

  vector<bool> seen(N, false);
  vi parent_road(N);
  map<int, set<int>> connect_road;

  rep(i, N) {
    if (seen[i])
      continue;
    seen[i] = true;
    queue<int> q;
    connect_road[i].insert(i);
    q.push(i);

    while (!q.empty()) {
      int v = q.front();
      q.pop();
      for (auto nv : road[v]) {
        if (seen[nv])
          continue;
        seen[nv] = true;
        connect_road[i].insert(nv);
        q.push(nv);
      }
    }

    for (auto v : connect_road[i]) {
      parent_road[v] = i;
    }
  }

  seen.assign(N, false);
  vi parent_train(N);
  map<int, set<int>> connect_train;

  rep(i, N) {
    if (seen[i])
      continue;
    seen[i] = true;
    queue<int> q;
    connect_train[i].insert(i);
    q.push(i);

    while (!q.empty()) {
      int v = q.front();
      q.pop();
      for (auto nv : train[v]) {
        if (seen[nv])
          continue;
        seen[nv] = true;
        connect_train[i].insert(nv);
        q.push(nv);
      }
    }

    for (auto v : connect_train[i]) {
      parent_train[v] = i;
    }
  }

  seen.assign(N, false);
  vi ans(N);
  rep(i, N) {
    if (seen[i])
      continue;
    int a = parent_road[i];
    int b = parent_train[i];
    vi connect;

    if (connect_road[a].size() < connect_train[b].size()) {
      for (auto v : connect_road[a]) {
        if (parent_train[i] == parent_train[v]) {
          connect.push_back(v);
          seen[v] = true;
        }
      }
    } else {
      for (auto v : connect_train[b]) {
        if (parent_road[i] == parent_road[v]) {
          connect.push_back(v);
          seen[v] = true;
        }
      }
    }

    for (auto v : connect) {
      ans[v] = connect.size();
    }
  }

  rep(i, N - 1) cout << ans[i] << ' ';
  cout << ans[N - 1] << endl;
  return 0;
}