Sunday, June 17, 2012

Union Find

I've coded an nice Union Find implementation that should be helpful to anyone, it's based on a 'Weighted union find' which also stores an array for the size of the 'leaves' if you imagine a group of union procedures as merges two sub-trees, and to make things better, it only require O(lg N) to both Union and Find procedures, good isn't it ? Here goes the code which is also available on github and hatena

#include 

using namespace std;

struct UnionFind {
    int N, *id, *sz;
    UnionFind(int _N) {
        N = _N;
        id = new int[_N];
        sz = new int[_N];
        for(int i = 0; i < _N; i++) {
            id[i] = i;
            sz[i] = 1;
        }
    }
    int root(int i) {
        while(i != id[i]) {
            id[i] = id[id[i]];
            i = id[i];
        }
        return i;
    }
    bool find(int p, int q) {
        return root(p) == root(q);
    }
    void unite(int p, int q) {
        int i = root(p);
        int j = root(q);
        if(sz[i] < sz[j]) {
            id[i] = j;
            sz[j] += sz[i];
        } else {
            id[j] = i;
            sz[i] += sz[j];
        }
    }
}