Optimize Union Find (Disjoint Set Union) implementation

Question

I have implemented a Disjoint Set Union in C++ for a Kattis problem. However, my solution is inefficient, as I get TLA (Time Limit Exceeded) for the first hidden test case. Can anyone spot an inefficiency with my code?

I have based my implementation on this article. In particular, I set the parent of the parent of the smaller set to the parent of the bigger set (when making an union of the sets). According to the mentioned article, this should be as efficient as using the rank:

Both optimizations are equivalent in terms of time and space complexity. So in practice you can use any of them.

#include <bits/stdc++.h>

using namespace std;

typedef vector<int> vi;

class UnionFind {
private:
  vi p;
  vi size;

public:
  UnionFind(int N) {
    p.assign(N, 0);
    for (int i = 0; i != N; ++i)
      p[i] = i;
    size.assign(N, 1);
  }

  int findSet(int i) { return (p[i] == i) ? i : (p[i] = findSet(p[i])); }

  bool isSameSet(int i, int j) { return findSet(i) == findSet(j); }

  void unionSet(int i, int j) {
    int x = findSet(i);
    int y = findSet(j);
    if (x == y)
      return;
    if (size[x] < size[y])
      swap(x, y);
    p[y] = x;
    size[x] += size[y];
  }
};

int main() {
  int n;
  int q;
  cin >> n >> q;
  auto set = UnionFind(n);
  while (q--) {
    char ch;
    cin >> ch;
    int x;
    int y;
    cin >> x >> y;
    if (ch == '=') {
      set.unionSet(x, y);
    } else if (ch == '?') {
      if (set.isSameSet(x, y))
        cout << "yes" << endl;
      else
        cout << "no" << endl;
    }
  }
}

Answer 1

This seems to be a Kattis-specific IO problem. Adding ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL); to the start of main, and changing both occurrences of endl to '\n' causes it to pass in .2 seconds.

I'll also mention that including <bits/stdc++.h> is frowned upon here, but the only cause of your error was the IO. Your implementation of Union Find is correct and runs in the standard inverse Ackermann time bounds.