MRUnionFind.h

Go to the documentation of this file.

#pragma once
 
#include "MRVector.h"
#include <utility>
 
namespace MR
{
 
template <typename I>
class UnionFind
{
public:
    UnionFind() = default;
    explicit UnionFind( size_t size ) { reset( size ); }
    auto size() const { return roots_.size(); }
 
    void reset( size_t size )
    {
        roots_.clear();
        roots_.reserve( size );
        for ( I i{ size_t( 0 ) }; i < size; ++i )
            roots_.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] )
        {
            roots_[firstRoot] = secondRoot;
            sizes_[secondRoot] += sizes_[firstRoot];
            return { secondRoot, true };
        }
        else
        {
            roots_[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;
    }
 
    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 < roots_.size(); ++i )
            updateRoot_( i, findRootNoUpdate_( i ) );
        return roots_;
    }
 
    const Vector<I, I> & parents() const { return roots_; }
 
    size_t sizeOfComp( I a ) { return sizes_[ find( a ) ]; }
 
private:
    I findRootNoUpdate_( I a ) const
    {
        I r = roots_[a];
        for ( I e = a; e != r; r = roots_[e = r] ) {}
        return r;
    }
    I updateRoot_( I a, const I r )
    {
        // update parents
        while ( a != r ) 
        { 
            I b = r;
            std::swap( roots_[a], b );
            a = b;
        }
        return r;
    }
    // find the root of given element, and set it as parent for it and other parents
    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( roots_[a], b );
                a = b;
            }
            else
                a = roots_[a];
        }
        return r;
    }
    Vector<I, I> roots_;
    Vector<size_t, I> sizes_;
};
 
}