Problem p03857 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define MAX 200005

struct unionfind {
  int pa[MAX];
  int ra[MAX];
  void init() {
    for (int i = 0; i < MAX; i++) {
      pa[i] = i;
      ra[i] = 0;
    }
  }

  int find(int x) {
    if (pa[x] == x)
      return x;
    else
      return pa[x] = find(pa[x]);
  }

  void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return;
    if (ra[x] < ra[y])
      swap(x, y);
    pa[y] = x;
    if (ra[x] == ra[y])
      ra[x]++;
  }

  bool same(int x, int y) { return (find(x) == find(y)); }
};

unionfind A, B;
set<int> s[MAX];
set<int> t[MAX];
int ans[MAX];

int N, K, L;

void solve(int id) {

  int x = A.find(id);
  int y = B.find(id);

  vector<int> v;
  for (int key : t[y]) {
    if (s[x].count(key) > 0) {
      v.push_back(key);
    }
  }

  for (int key : v) {
    ans[key] = v.size();
    if (t[x].count(key))
      t[x].erase(key);
    if (t[y].count(key))
      t[y].erase(key);
  }
}

int main() {
  scanf("%d %d %d", &N, &K, &L);
  A.init();
  B.init();
  for (int i = 0; i < K; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    A.unite(a, b);
  }
  for (int i = 0; i < L; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    B.unite(a, b);
  }

  for (int i = 1; i <= N; i++) {
    s[A.find(i)].insert(i);
    t[B.find(i)].insert(i);
  }

  memset(ans, -1, sizeof(ans));
  for (int i = 1; i <= N; i++) {
    if (ans[i] == -1) {
      solve(i);
    }
  }
  for (int i = 1; i <= N; i++) {
    if (i > 1)
      printf(" ");
    printf("%d", ans[i]);
  }
  printf("\n");
  return 0;
}
// 123 4
// 14 23

// 1235 67 4
// 1 2 345 67

/*
for(int i=1;i<=N;i++){
  cout<<i<<endl;
  for(int x:s[i]){
    cout<<x<<' ';
  }
  cout<<endl;
}
for(int i=1;i<=N;i++){
  cout<<i<<endl;
  for(int x:t[i]){
    cout<<x<<' ';
  }
  cout<<endl;
}
*/


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define MAX 200005

struct unionfind {
  int pa[MAX];
  int ra[MAX];
  void init() {
    for (int i = 0; i < MAX; i++) {
      pa[i] = i;
      ra[i] = 0;
    }
  }

  int find(int x) {
    if (pa[x] == x)
      return x;
    else
      return pa[x] = find(pa[x]);
  }

  void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return;
    if (ra[x] < ra[y])
      swap(x, y);
    pa[y] = x;
    if (ra[x] == ra[y])
      ra[x]++;
  }

  bool same(int x, int y) { return (find(x) == find(y)); }
};

unionfind A, B;
set<int> s[MAX];
set<int> t[MAX];
int ans[MAX];

int N, K, L;

void solve(int id) {

  int x = A.find(id);
  int y = B.find(id);

  vector<int> v;
  if (t[y].size() < s[x].size()) {
    for (int key : t[y]) {
      if (s[x].count(key) > 0) {
        v.push_back(key);
      }
    }
  } else {
    for (int key : s[x]) {
      if (t[y].count(key) > 0) {
        v.push_back(key);
      }
    }
  }

  for (int key : v) {
    ans[key] = v.size();
    if (t[x].count(key))
      t[x].erase(key);
    if (t[y].count(key))
      t[y].erase(key);
  }
}

int main() {
  scanf("%d %d %d", &N, &K, &L);
  A.init();
  B.init();
  for (int i = 0; i < K; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    A.unite(a, b);
  }
  for (int i = 0; i < L; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    B.unite(a, b);
  }

  for (int i = 1; i <= N; i++) {
    s[A.find(i)].insert(i);
    t[B.find(i)].insert(i);
  }

  memset(ans, -1, sizeof(ans));
  for (int i = 1; i <= N; i++) {
    if (ans[i] == -1) {
      solve(i);
    }
  }
  for (int i = 1; i <= N; i++) {
    if (i > 1)
      printf(" ");
    printf("%d", ans[i]);
  }
  printf("\n");
  return 0;
}
// 123 4
// 14 23

// 1235 67 4
// 1 2 345 67

/*
for(int i=1;i<=N;i++){
  cout<<i<<endl;
  for(int x:s[i]){
    cout<<x<<' ';
  }
  cout<<endl;
}
for(int i=1;i<=N;i++){
  cout<<i<<endl;
  for(int x:t[i]){
    cout<<x<<' ';
  }
  cout<<endl;
}
*/