Problem p03356 (C++)

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

#define rep(i, n) for (int i = 0; i < (n); ++i)

typedef long long ll;
typedef pair<int, int> P;

const int MAX = 1e5 + 5;
const int INF = 1001001001;
const int MOD = 1e9 + 7;

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

class UnionFindTree {
private:
  vector<int> par, rank;

public:
  UnionFindTree(int n) {
    par = rank = vector<int>(n);

    for (int i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

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

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    if (rank[x] < rank[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }

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

int main() {
  int N, M;
  cin >> N >> M;

  vector<int> p(N);
  rep(i, N) {
    cin >> p[i];
    p[i]--;
  }

  vector<int> X(M), Y(M);
  rep(i, M) {
    int x, y;
    cin >> x >> y;
    x--, y--;
    X[i] = x, Y[i] = y;
  }

  UnionFindTree uf = UnionFindTree(N);
  rep(i, M) { uf.unite(X[i], Y[i]); }

  map<int, vector<int>> m;
  rep(i, N) {
    int p = uf.find(i);
    m[p].push_back(i);
  }

  int ans = 0;
  for (auto &i : m) {
    vector<int> &v = i.second;
    vector<int> cnt(N);
    for (int j : v) {
      cnt[j]++;
      cnt[p[j]]++;
      if (cnt[j] == 2)
        ans++;
      if (j != p[j] && cnt[p[j]] == 2)
        ans++;
    }
  }

  cout << ans << endl;
}

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

#define rep(i, n) for (int i = 0; i < (n); ++i)

typedef long long ll;
typedef pair<int, int> P;

const int MAX = 1e5 + 5;
const int INF = 1001001001;
const int MOD = 1e9 + 7;

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

class UnionFindTree {
private:
  vector<int> par, rank;

public:
  UnionFindTree(int n) {
    par = rank = vector<int>(n);

    for (int i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

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

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    if (rank[x] < rank[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }

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

int main() {
  int N, M;
  cin >> N >> M;

  vector<int> p(N);
  rep(i, N) {
    cin >> p[i];
    p[i]--;
  }

  vector<int> X(M), Y(M);
  rep(i, M) {
    int x, y;
    cin >> x >> y;
    x--, y--;
    X[i] = x, Y[i] = y;
  }

  UnionFindTree uf = UnionFindTree(N);
  rep(i, M) { uf.unite(X[i], Y[i]); }

  map<int, vector<int>> m;
  rep(i, N) {
    int p = uf.find(i);
    m[p].push_back(i);
  }

  int ans = 0;
  for (auto &i : m) {
    vector<int> &v = i.second;
    map<int, int> cnt;
    for (int j : v) {
      cnt[j]++;
      cnt[p[j]]++;
      if (cnt[j] == 2)
        ans++;
      if (j != p[j] && cnt[p[j]] == 2)
        ans++;
    }
  }

  cout << ans << endl;
}