Implementing path compression in a disjoint set data structure?

Question

Here is my Objective-C implementation of a disjoint set. - Positive number point to parent. - Negative number indicate root & children count. (So they each start disjointed at -1.) - The index acts as the data I am grouping. Seems to work ok... just had a couple questions.

1. find: How can I compress the path? Because I am not doing it recursively, do I have to store the path and loop it again to set after find root?

2. join: I am basing join on children count instead of depth!? I guess that is not right. Do I need to do something special during join if depths equal?

Thanks.

DisjointSet.h

@interface DisjointSet : NSObject
{
    NSMutableArray *_array;
}

- (id)initWithSize:(NSInteger)size;
- (NSInteger)find:(NSInteger)item;
- (void)join:(NSInteger)root1 root2:(NSInteger)root2;

@end

DisjointSet.m

#import "DisjointSet.h"

@implementation DisjointSet

- (id)initWithSize:(NSInteger)size
{
    self = [super init];
    if (self)
    {
        _array = [NSMutableArray arrayWithCapacity:size];
        for (NSInteger i = 0; i < size; i++)
        {
            [_array addObject:[NSNumber numberWithInteger:-1]];
        }
    }
    return self;
}

- (NSInteger)find:(NSInteger)item
{
    while ([[_array objectAtIndex:item] integerValue] >= 0)
    {
        item = [[_array objectAtIndex:item] integerValue];
    }
    return item;
}

- (void)join:(NSInteger)root1 root2:(NSInteger)root2
{
    if (root1 == root2) return;

    NSInteger data1 = [[_array objectAtIndex:root1] integerValue];
    NSInteger data2 = [[_array objectAtIndex:root2] integerValue];
    if (data2 < data1)
    {
        [_array setObject:[NSNumber numberWithInteger:data2 + data1] atIndexedSubscript:root2];
        [_array setObject:[NSNumber numberWithInteger:root2] atIndexedSubscript:root1];
    }
    else
    {
        [_array setObject:[NSNumber numberWithInteger:data1 + data2] atIndexedSubscript:root1];
        [_array setObject:[NSNumber numberWithInteger:root1] atIndexedSubscript:root2];
    }
}

@end

Answer 1

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