# Union Find implementation

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
}
return 0;
}


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];
}
}


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);

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.