For the find operation, there is no need to store the path (separately from your _array
) or to use recursion. Either of those approaches requires O(P) storage (P = path length). Instead, you can just traverse the path twice. The first time, you find the root. The second time, you set all of the children to point to the root. This takes O(P) time and O(1) storage.
- (NSInteger)findItem:(NSInteger)item {
NSInteger root;
NSNumber *rootObject = nil;
for (NSInteger i = item; !rootObject; ) {
NSInteger parent = [_array[i] integerValue];
if (parent < 0) {
root = i;
rootObject = @(i);
}
i = parent;
}
for (NSInteger i = item; i != root; ) {
NSInteger parent = [_array[i] integerValue];
_array[i] = rootObject;
i = parent;
}
return root;
}
For the merge operation, you want to store each root's rank (which is an upper bound on its depth), not each root's descendant count. Storing each root's rank allows you to merge the shorter tree into the taller tree, which guarantees O(log N) time for find operations. The rank only increases when the trees to be merged have equal rank.
- (void)joinItem:(NSInteger)a item:(NSInteger)b {
NSInteger aRank = -[_array[a] integerValue];
NSInteger bRank = -[_array[b] integerValue];
if (aRank < bRank) {
NSInteger t = a;
a = b;
b = t;
} else if (aRank == bRank) {
_array[a] = @(-aRank - 1);
}
_array[b] = @(a);
}