MRUnionFind.h

Go to the documentation of this file.

#pragma once
 
#include "MRVector.h"
#include "MRTimer.h"
#include <utility>
 
namespace MR
{
 
template <typename I>
class UnionFind
{
public:
    using SizeType = typename I::ValueType;
 
    UnionFind() = default;
 
    explicit UnionFind( size_t size ) { reset( size ); }
 
    auto size() const { return parents_.size(); }
 
    void reset( size_t size )
    {
        MR_TIMER;
        parents_.clear();
        parents_.reserve( size );
        for ( I i{ size_t( 0 ) }; i < size; ++i )
            parents_.push_back( i );
        sizes_.clear();
        sizes_.resize( size, 1 );
    }
 
    std::pair<I,bool> unite( I first, I second )
    {
        auto firstRoot = updateRoot_( first );
        auto secondRoot = updateRoot_( second );
        if ( firstRoot == secondRoot )
            return { firstRoot, false };
        if ( sizes_[firstRoot] < sizes_[secondRoot] )
        {
            parents_[firstRoot] = secondRoot;
            sizes_[secondRoot] += sizes_[firstRoot];
            return { secondRoot, true };
        }
        else
        {
            parents_[secondRoot] = firstRoot;
            sizes_[firstRoot] += sizes_[secondRoot];
            return { firstRoot, true };
        }
    }
 
    bool united( I first, I second )
    {
        auto firstRoot = updateRoot_( first );
        auto secondRoot = updateRoot_( second );
        return firstRoot == secondRoot;
    }
 
    bool isRoot( I a ) const { return parents_[a] == a; }
 
    I parent( I a ) const { return parents_[a]; }
 
    I find( I a ) { return updateRoot_( a, findRootNoUpdate_( a ) ); }
 
    I findUpdateRange( I a, I begin, I end )
    {
        return updateRootInRange_( a, findRootNoUpdate_( a ), begin, end );
    }
 
    const Vector<I, I> & roots()
    {
        for ( I i{ size_t( 0 ) }; i < parents_.size(); ++i )
            updateRoot_( i, findRootNoUpdate_( i ) );
        return parents_;
    }
 
    const Vector<I, I> & parents() const { return parents_; }
 
    SizeType sizeOfComp( I a ) { return sizes_[ find( a ) ]; }
 
private:
    I findRootNoUpdate_( I a ) const
    {
        I r = parents_[a];
        for ( I e = a; e != r; r = parents_[e = r] ) {}
        return r;
    }
 
    I updateRoot_( I a, const I r )
    {
        // update parents
        while ( a != r ) 
        { 
            I b = r;
            std::swap( parents_[a], b );
            a = b;
        }
        return r;
    }
 
    I updateRoot_( I a ) { return updateRoot_( a, findRootNoUpdate_( a ) ); }
 
    I updateRootInRange_( I a, const I r, I begin, I end )
    {
        assert( begin < end );
        // update parents
        while ( a != r )
        { 
            if ( a >= begin && a < end )
            {
                I b = r;
                std::swap( parents_[a], b );
                a = b;
            }
            else
                a = parents_[a];
        }
        return r;
    }
 
    Vector<I, I> parents_;
 
    Vector<SizeType, I> sizes_;
};
 
} //namespace MR