MRUnionFind.h

Go to the documentation of this file.

#pragma once
 
#include "MRVector.h"
#include "MRTimer.h"
#include <atomic>
#include <utility>
 
namespace MR
{
 

template <typename I>
class BaseUnionFind
{
public:
    using SizeType = typename I::ValueType;
 
    BaseUnionFind() = default;
    BaseUnionFind( const BaseUnionFind & ) = default;
    BaseUnionFind( BaseUnionFind && ) noexcept = default;
    BaseUnionFind & operator =( const BaseUnionFind & ) = default;
    BaseUnionFind & operator =( BaseUnionFind && ) noexcept = default;

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

    std::pair<I, bool> uniteUnbalanced( I first, I second )
    {
        auto firstRoot = updateRoot_( first );
        auto secondRoot = updateRoot_( second );
        if ( firstRoot == secondRoot )
            return { firstRoot, false };
        parents_[firstRoot] = secondRoot;
        return { secondRoot, true };
    }

    bool united( I first, I second )
    {
        return updateRoot_( first ) == updateRoot_( second );
    }

    bool isRoot( I a ) const { return parents_[a] == a; }

    I parent( I a ) const { return parents_[a]; }

    I find( I a ) { return updateRoot_( 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_; }
 
protected:
    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 )
    {
        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 );
        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_;
};

template <typename I>
class UnionFind : public BaseUnionFind<I>
{
public:
    using SizeType = typename I::ValueType;
 
    UnionFind() = default;

    explicit UnionFind( size_t size ) { reset( size ); }

    explicit UnionFind( BaseUnionFind<I> && base ) : BaseUnionFind<I>( std::move( base ) )
    {
        sizes_.resize( this->parents_.size(), SizeType( 0 ) );
        for ( I i{ size_t( 0 ) }; i < this->parents_.size(); ++i )
            ++sizes_[ this->findRootNoUpdate_( i ) ];
    }

    void reset( size_t size )
    {
        BaseUnionFind<I>::reset( size );
        sizes_.clear();
        sizes_.resize( size, 1 );
    }

    std::pair<I,bool> unite( I first, I second )
    {
        auto firstRoot = this->updateRoot_( first );
        auto secondRoot = this->updateRoot_( second );
        if ( firstRoot == secondRoot )
            return { firstRoot, false };
        if ( sizes_[firstRoot] < sizes_[secondRoot] )
        {
            this->parents_[firstRoot] = secondRoot;
            sizes_[secondRoot] += sizes_[firstRoot];
            return { secondRoot, true };
        }
        else
        {
            this->parents_[secondRoot] = firstRoot;
            sizes_[firstRoot] += sizes_[secondRoot];
            return { firstRoot, true };
        }
    }

    SizeType sizeOfComp( I a ) { return sizes_[ this->find( a ) ]; }
 
private:
    Vector<SizeType, I> sizes_;
};

template <typename I>
class ParallelUnionFind : public BaseUnionFind<I>
{
public:
    ParallelUnionFind() = default;

    explicit ParallelUnionFind( size_t size ) { this->reset( size ); }

    void uniteAtomic( I first, I second )
    {
        for ( ;; )
        {
            first = findAtomic_( first );
            second = findAtomic_( second );
            if ( first == second )
                return; 
            if ( first < second )
                std::swap( first, second );
            if ( casStrong_( this->parents_[first], first, second ) )
                return;
        }
    }
 
private:

    static I loadAtomic_( I& slot )
    {
#ifdef __cpp_lib_atomic_ref
        return I( std::atomic_ref<typename I::ValueType>( slot.get() ).load( std::memory_order_relaxed ) );
#else
        return I( __atomic_load_n( &slot.get(), __ATOMIC_RELAXED ) );
#endif
    }

    static void casWeak_( I& slot, I expected, I desired )
    {
        using T = typename I::ValueType;
        T exp = expected.get();
#ifdef __cpp_lib_atomic_ref
        std::atomic_ref<T>( slot.get() ).compare_exchange_weak( exp, desired.get(), std::memory_order_relaxed );
#else
        __atomic_compare_exchange_n( &slot.get(), &exp, desired.get(), true, __ATOMIC_RELAXED, __ATOMIC_RELAXED );
#endif
    }

    static bool casStrong_( I& slot, I expected, I desired )
    {
        using T = typename I::ValueType;
        T exp = expected.get();
#ifdef __cpp_lib_atomic_ref
        return std::atomic_ref<T>( slot.get() ).compare_exchange_strong( exp, desired.get(), std::memory_order_acq_rel, std::memory_order_relaxed );
#else
        return __atomic_compare_exchange_n( &slot.get(), &exp, desired.get(), false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED );
#endif
    }

    I findAtomic_( I a )
    {
        for ( ;; )
        {
            I p = loadAtomic_( this->parents_[a] );
            if ( p == a )
                return a; 
            I gp = loadAtomic_( this->parents_[p] );
            if ( p == gp )
                return p; 
            casWeak_( this->parents_[a], p, gp );
            a = gp;
        }
    }
};
 
}