Union Find implementation

Question

I am trying to complete this challenge. The user should enter a sequence of instructions, = to link two numbers, and ? to query whether two integers are linked. For example,

? 1 2
= 1 5
= 2 5
? 1 2

should produce

no
yes

I have tried following the algorithm given here to the letter, but I keep getting a judgment of "time limit exceeded". It's very infuriating as the difficultly rating for the problem suggests it should be quite easy. I have tried modifying the code so that redundant links (e.g. = 1 2 where 1 is never queried nor later linked to another number) are ignored, but it still isn't fast enough.

#include <iostream>
#include <vector>

using namespace std;

int root(int a, vector<int> & parent) {
    int b = parent[a];
    return b == a ? a : parent[a] = root(b, parent);
}

void link(int a1, int a2, vector<int> & parent, vector<int> & rk) {
    int root1 = root(a1, parent);
    int root2 = root(a2, parent);
    if (root1 == root2)
        return;
    int rk1 = rk[root1];
    int rk2 = rk[root2];
    if (rk1 < rk2) {
        parent[root1] = root2;
    } else if (rk2 < rk1) {
        parent[root2] = root1;
    } else {
        parent[root1] = root2;
        rk[root2]++;
    }
}

int main() {
    int n, q;
    cin >> n >> q;
    vector<int> parent;
    for (int i = 0; i < n; i++)
        parent.push_back(i);
    vector<int> rk(n, 0);
    for (int i = 0; i < q; i++) {
        string str;
        int a, b;
        cin >> str >> a >> b;
        if (str == "?")
            cout << (a == b || root(a, parent) == root(b, parent) ? "yes" : "no") << endl;
        else
            link(a, b, parent, rk);
    }
    return 0;
}

Answer 1

Consider an iterative path compression design for root. The recursive version uses more memory for stack frames. Typical iterative versions follow the path one at a time. The path can be followed two at a time since the end loops back to itself. The root function could be refactored to be something along the lines of:

int root(int a, vector<int> &parent) {
    while (parent[a] != a)
      a = parent[a] = parent[parent[a]];
    return a;
}

Plus, consider adding a second vector to keep track of the sizes of each group. This will allow for more intelligent selecting of representatives during link. Something along the lines of the following:

void link(int i, int j, vector<int> &parent, vector<int> &sizes) {
  i = root(i, parent);
  j = root(j, parent);
  if (i == j) return;
  if (sizes[i] < sizes[j]) {
    parent[i] = parent[j];
    sizes[j] += sizes[i];
  } else {
    parent[j] = parent[i];
    sizes[i] += sizes[j];
  }
}

See this post for more information.

Consider allocating the vector using the constructor, resize, or reserve. Then initial values can be set using iota (ref). For example,

vector<int> parent(n);
iota(begin(parent), end(parent), 0);

Consider using printf and scanf over cin and cout. Also, if you are going to use cin and cout, consider including ios_base::sync_with_stdio(false); and cin.tie(NULL);. See this question for more information.