DisjointSets.h

/*
 * Copyright 2022 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
#pragma once
 
#include <lagrange/utils/span.h>
 
#include <vector>
 
namespace lagrange {
 
 
template <typename IndexType>
class DisjointSets
{
public:
    DisjointSets() = default;
 
    explicit DisjointSets(size_t n) { init(n); }
 
    void init(size_t n);
 
    size_t size() const { return m_parent.size(); }
 
    void clear() { m_parent.clear(); }
 
    IndexType find(IndexType i);
 
    IndexType merge(IndexType i, IndexType j)
    {
        const auto root_i = find(i);
        const auto root_j = find(j);
        return m_parent[root_j] = root_i;
    }
 
    [[deprecated]] std::vector<std::vector<IndexType>> extract_disjoint_sets();
 
    size_t extract_disjoint_set_indices(std::vector<IndexType>& index_map);
 
    size_t extract_disjoint_set_indices(span<IndexType> index_map);
 
protected:
    std::vector<IndexType> m_parent;
};
 
template <typename IndexType>
class DisjointSetsWithSize : public DisjointSets<IndexType>
{
    using Super = DisjointSets<IndexType>;
 
public:
    using DisjointSets<IndexType>::DisjointSets;
 
    explicit DisjointSetsWithSize(size_t n) { init(n); }
 
    void init(size_t n)
    {
        Super::init(n);
        m_size.assign(n, 1);
    }
 
    void clear()
    {
        Super::clear();
        m_size.clear();
    }
 
    IndexType size_of(IndexType i) { return m_size[this->find(i)]; }
 
    IndexType merge(IndexType i, IndexType j)
    {
        auto root_i = this->find(i);
        auto root_j = this->find(j);
 
        if (root_i == root_j) {
            return root_i;
        }
 
        // Ensure |I| >= |J|
        if (size_of(root_i) < size_of(root_j)) {
            std::swap(root_i, root_j);
        }
 
        this->m_parent[root_j] = root_i;
        m_size[root_i] += m_size[root_j];
 
        return root_i;
    }
 
protected:
    std::vector<IndexType> m_size;
};
 
 
} // namespace lagrange