bucket_sort.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/internal/invert_mapping.h>
#include <lagrange/utils/DisjointSets.h>
 
#include <numeric>
 
namespace lagrange::internal {
 
template <typename Index>
struct BucketSortResult
{
    Index num_representatives;
 
    std::vector<Index> sorted_elements;
 
    std::vector<Index> representative_offsets;
};
 
template <typename Index>
BucketSortResult<Index> bucket_sort(
    DisjointSets<Index>& unified_indices,
    span<Index> element_representative)
{
    la_debug_assert(unified_indices.size() == element_representative.size());
    BucketSortResult<Index> result;
 
    auto& num_representatives = result.num_representatives;
    auto& sorted_elements = result.sorted_elements;
    auto& representative_offsets = result.representative_offsets;
 
    Index num_elements = static_cast<Index>(element_representative.size());
 
    // Calculate representative element for each bucket
    num_representatives = 0;
    std::fill(element_representative.begin(), element_representative.end(), invalid<Index>());
    for (Index e = 0; e < num_elements; ++e) {
        Index r = unified_indices.find(e);
        if (element_representative[r] == invalid<Index>()) {
            element_representative[r] = num_representatives++;
        }
        element_representative[e] = element_representative[r];
    }
 
    auto mapping = invert_mapping(
        {element_representative.data(), element_representative.size()},
        num_representatives);
    sorted_elements = std::move(mapping.data);
    representative_offsets = std::move(mapping.offsets);
 
    return result;
}
 
template <typename Index>
struct BucketSortOffset
{
    Index num_representatives;
 
    std::vector<Index> representative_offsets;
};
 
template <typename Index, typename Function>
BucketSortOffset<Index>
bucket_sort(std::vector<Index>& elements, Index num_buckets, Function get_representative)
{
    BucketSortOffset<Index> result;
 
    auto& num_representatives = result.num_representatives;
    auto& representative_offsets = result.representative_offsets;
    num_representatives = num_buckets;
    auto mapping = invert_mapping<Index>(
        static_cast<Index>(elements.size()),
        get_representative,
        num_representatives);
    elements = std::move(mapping.data);
    representative_offsets = std::move(mapping.offsets);
    return result;
}
 
} // namespace lagrange::internal