DisjointSetUnionFind.h

#ifndef PBAT_GRAPH_DISJOINTSETUNIONFIND_H
#define PBAT_GRAPH_DISJOINTSETUNIONFIND_H
 
#include "pbat/Aliases.h"
#include "pbat/common/Concepts.h"
 
#include <cassert>
 
namespace pbat::graph {

template <common::CIndex TIndex = Index>
class DisjointSetUnionFind
{
  public:
    using IndexType = TIndex; 
 
    DisjointSetUnionFind() = default;
    DisjointSetUnionFind(IndexType n);
    void Prepare(IndexType n);
    IndexType Find(IndexType u);
    IndexType Union(IndexType u, IndexType v);
    IndexType Root(IndexType u) const;
    IndexType Size(IndexType u) const;
    IndexType NumVertices() const;
 
  private:
    Eigen::Vector<IndexType, Eigen::Dynamic> mParent;
    Eigen::Vector<IndexType, Eigen::Dynamic> mRank;
    Eigen::Vector<IndexType, Eigen::Dynamic> mSize;
};
 
template <common::CIndex TIndex>
DisjointSetUnionFind<TIndex>::DisjointSetUnionFind(IndexType n) : DisjointSetUnionFind<TIndex>()
{
    Prepare(n);
}
 
template <common::CIndex TIndex>
void DisjointSetUnionFind<TIndex>::Prepare(IndexType n)
{
    mParent = Eigen::Vector<IndexType, Eigen::Dynamic>::LinSpaced(n, IndexType(0), n - 1);
    mRank.setZero(n);
    mSize.setOnes(n);
}
 
template <common::CIndex TIndex>
TIndex DisjointSetUnionFind<TIndex>::Find(IndexType u)
{
    if (mParent[u] != u)
        mParent[u] = Find(mParent[u]); // Path compression
    return mParent[u];
}
 
template <common::CIndex TIndex>
TIndex DisjointSetUnionFind<TIndex>::Union(IndexType u, IndexType v)
{
    assert(u != v && "Cannot union the same vertex with itself.");
    // Find roots of each node
    IndexType const ru = Find(u);
    IndexType const rv = Find(v);
    // Make tree with smaller height a subtree of the other tree.
    // This effectively minimizes the depth of the resulting merged tree.
    if (mRank[ru] > mRank[rv])
    {
        mParent[rv] = ru;
        mSize[ru] += mSize[rv];
        mSize[rv] = IndexType(0);
    }
    else
    {
        mParent[ru] = rv;
        mSize[rv] += mSize[ru];
        mSize[ru] = IndexType(0);
    }
    // If both trees had the same depth, then
    // making one tree a subtree of the other increased
    // the depth of the merged tree by 1.
    if (mRank[ru] == mRank[rv])
    {
        ++mRank[rv];
    }
    return (mParent[ru] == ru) ? ru : rv;
}
 
template <common::CIndex TIndex>
inline TIndex DisjointSetUnionFind<TIndex>::Root(IndexType u) const
{
    while (mParent[u] != u)
        u = mParent[u];
    return u;
}
 
template <common::CIndex TIndex>
TIndex DisjointSetUnionFind<TIndex>::Size(IndexType u) const
{
    return mSize[Root(u)];
}
 
template <common::CIndex TIndex>
TIndex DisjointSetUnionFind<TIndex>::NumVertices() const
{
    return static_cast<TIndex>(mSize.size());
}
 
} // namespace pbat::graph
 
#endif // PBAT_GRAPH_DISJOINTSETUNIONFIND_H