union_find_set.cpp

#include <iostream>
#include <vector>

// rank based union, with path compressed
class UnionFindSet
{
public:
    UnionFindSet(const int N); 

    void Union(int p, int q); 
    int Find(int p); 
    bool Connected(int p, int q); 

private:
    std::vector<int> id; 
    std::vector<int> rank;
};

UnionFindSet::UnionFindSet(const int N)
    : rank(N, 0)
{
    id.reserve(N);
    for (int i=0; i<N; ++i) {
        id[i] = i;
    }   
}

void UnionFindSet::Union(int p, int q)
{
    int p_id = Find(p), q_id = Find(q);
    if (p_id != q_id) {
        if (rank[p_id] <= rank[q_id]) {
            id[p_id] = q_id;
            if (rank[p_id] == rank[q_id]) {
                ++rank[q_id];
            }   
        } else {
            id[q_id] = p_id;
        }
    }
}

int UnionFindSet::Find(int p)
{
    int root;
    for(root=p; root!=id[root]; root=id[root]);
    int q;
    while(p != id[p]) {
        q = id[p];
        id[p] = root;
        p = q;
    }
    return root;
}

inline bool UnionFindSet::Connected(int p, int q)
{
    return Find(p) == Find(q);
}

int main(int argc, const char *argv[])
{
    UnionFindSet ufs(10);
    ufs.Union(1, 3);
    ufs.Union(2, 4);
    std::cout << ufs.Connected(1, 3) << std::endl;
    std::cout << ufs.Connected(2, 4) << std::endl;
    ufs.Union(1, 4);
    std::cout << ufs.Connected(2, 3) << std::endl;
    return 0;
}