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.