Problem p03356 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cmath>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

template <class T> class UnionFind {

private:
  struct subset {
    T parent;
    T rank;
    T diff_weight;
    T sizes;
  };
  std::vector<subset> subsets;

public:
  UnionFind(T N) {
    subsets.reserve(N);
    for (T i = 0; i < N; i++) {
      subsets.push_back({i, 0, 0, 1});
    }
  }

  T Find(T x) {
    if (subsets[x].parent == x)
      return x;
    subsets[x].diff_weight += subsets[subsets[x].parent].diff_weight;
    return subsets[x].parent = Find(subsets[x].parent);
  }

  T weight(T x) {
    Find(x);
    return subsets[x].diff_weight;
  }

  bool Same(T x, T y) { return Find(x) == Find(y); }

  bool Union(T x, T y, T w) {
    w += weight(x);
    w -= weight(y);

    T xroot = Find(x);
    T yroot = Find(y);
    if (xroot == yroot)
      return false;

    if (subsets[xroot].rank < subsets[yroot].rank) {
      subsets[xroot].parent = yroot;
      w = -w;
      subsets[yroot].sizes += subsets[xroot].sizes;
    } else if (subsets[xroot].rank > subsets[yroot].rank) {
      subsets[yroot].parent = xroot;
      subsets[xroot].sizes += subsets[yroot].sizes;
    } else {
      subsets[yroot].parent = xroot;
      subsets[xroot].rank++;
      subsets[xroot].sizes += subsets[yroot].sizes;
    }
    subsets[yroot].diff_weight = w;
    return true;
  }

  T diff(T x, T y) { return weight(y) - weight(x); }

  T size(T x) { return subsets[Find(x)].sizes; }
};

int main(void) {
  int N, M;
  std::cin >> N >> M;

  std::vector<int> p(N);
  for (int i = 0; i < N; i++)
    std::cin >> p[i];

  UnionFind<int> uf(N);
  for (int i = 0; i < M; i++) {
    int a, b;
    std::cin >> a >> b;
    uf.Union(a - 1, b - 1, 0);
  }

  std::unordered_map<int, std::vector<int>> P;
  for (int i = 0; i < N; i++) {
    P[uf.Find(i)].push_back(i);
  }

  int ans = 0;
  for (auto itr = P.begin(); itr != P.end(); ++itr) {
    for (auto v = itr->second.begin(); v != itr->second.end(); ++v) {
      if (std::find(itr->second.begin(), itr->second.end(), p[*v] - 1) !=
          itr->second.end())
        ans++;
    }
  }

  std::cout << ans << std::endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

template <class T> class UnionFind {

private:
  struct subset {
    T parent;
    T rank;
    T diff_weight;
    T sizes;
  };
  std::vector<subset> subsets;

public:
  UnionFind(T N) {
    subsets.reserve(N);
    for (T i = 0; i < N; i++) {
      subsets.push_back({i, 0, 0, 1});
    }
  }

  T Find(T x) {
    if (subsets[x].parent == x)
      return x;
    subsets[x].diff_weight += subsets[subsets[x].parent].diff_weight;
    return subsets[x].parent = Find(subsets[x].parent);
  }

  T weight(T x) {
    Find(x);
    return subsets[x].diff_weight;
  }

  bool Same(T x, T y) { return Find(x) == Find(y); }

  bool Union(T x, T y, T w) {
    w += weight(x);
    w -= weight(y);

    T xroot = Find(x);
    T yroot = Find(y);
    if (xroot == yroot)
      return false;

    if (subsets[xroot].rank < subsets[yroot].rank) {
      subsets[xroot].parent = yroot;
      w = -w;
      subsets[yroot].sizes += subsets[xroot].sizes;
    } else if (subsets[xroot].rank > subsets[yroot].rank) {
      subsets[yroot].parent = xroot;
      subsets[xroot].sizes += subsets[yroot].sizes;
    } else {
      subsets[yroot].parent = xroot;
      subsets[xroot].rank++;
      subsets[xroot].sizes += subsets[yroot].sizes;
    }
    subsets[yroot].diff_weight = w;
    return true;
  }

  T diff(T x, T y) { return weight(y) - weight(x); }

  T size(T x) { return subsets[Find(x)].sizes; }
};

int main(void) {
  int N, M;
  std::cin >> N >> M;

  std::vector<int> p(N);
  for (int i = 0; i < N; i++)
    std::cin >> p[i];

  UnionFind<int> uf(N);
  for (int i = 0; i < M; i++) {
    int a, b;
    std::cin >> a >> b;
    uf.Union(a - 1, b - 1, 0);
  }

  std::unordered_map<int, std::vector<int>> P;
  for (int i = 0; i < N; i++) {
    P[uf.Find(i)].push_back(i);
  }

  int ans = 0;
  for (auto itr = P.begin(); itr != P.end(); ++itr) {
    for (auto v = itr->second.begin(); v != itr->second.end(); ++v) {
      if (std::binary_search(itr->second.begin(), itr->second.end(), p[*v] - 1))
        ans++;
    }
  }

  std::cout << ans << std::endl;
  return 0;
}