random_data_sampling.h

#ifndef CPP_ALGORITHM_RANDOM_DATA_SAMPLING_H
#define CPP_ALGORITHM_RANDOM_DATA_SAMPLING_H

#include <numeric>
#include <random>
#include <vector>

namespace RandomDataSampling
{
    /**
     * \brief Select the k elements randomly from the array with uniform probability.
     * Let A be an array with n distinct elements. Design an algorithm to return a random subset of k elements from A,
     * where every subset has an equal chance of being chosen.
     * \param k sample size
     * \param arr input array
     * \return result array
     */
    std::vector<int> OfflineRandomSampling(
        int k,
        std::vector<int>& arr);

    /**
     * \brief Select the k elements randomly from the array with uniform probability.
     * Design an algorithm that reads data and creates a random subset of size k,
     * where each item has an equal chance of being included.
     * \param begin begin iterator
     * \param end end iterator
     * \param k sample size
     * \return result array
     */
    std::vector<int> OnlineRandomSampling(
        std::vector<int>::const_iterator begin,
        std::vector<int>::const_iterator end,
        int k);

    /**
     * \brief Compute permutation of the array generated by random sampling.
     * \param n sample size
     * \return result array
     */
    std::vector<int> ComputeRandomPermutation(int n);
}

// ----------------------------------------------------------------------------
inline std::vector<int> RandomDataSampling::OfflineRandomSampling(
    const int k,
    std::vector<int>& arr)
{
    std::default_random_engine seed((std::random_device())());
    for (int i = 0; i < k; ++i)
    {
        std::uniform_int_distribution<int> dist(i, static_cast<int>(arr.size()) - 1);
        std::swap(arr[i], arr[dist(seed)]);
    }
    return std::vector<int>{arr.begin(), arr.begin() + k};
}

// ----------------------------------------------------------------------------
inline std::vector<int> RandomDataSampling::OnlineRandomSampling(
    std::vector<int>::const_iterator begin,
    const std::vector<int>::const_iterator end,
    const int k)
{
    std::vector<int> running_sample;
    // save the first k elements
    for (int i = 0; i < k; ++i)
    {
        running_sample.emplace_back(*begin++);
    }

    std::default_random_engine seed((std::random_device())());
    int num_seen_so_far = k;
    while (begin != end)
    {
        int x = *begin++;
        ++num_seen_so_far;

        // generate a random number in [0, num_seen_so_far].
        // if the generated number exists in [0, k), replace the element with x.
        const int idx_to_replace = std::uniform_int_distribution<int>{0, num_seen_so_far - 1}(seed);
        if (idx_to_replace < k)
        {
            running_sample[idx_to_replace] = x;
        }
    }

    return running_sample;
}

// ----------------------------------------------------------------------------
inline std::vector<int> RandomDataSampling::ComputeRandomPermutation(const int n)
{
    auto permutation = std::vector<int>(n);
    std::iota(permutation.begin(), permutation.end(), 0);
    return OfflineRandomSampling(n, permutation);
}

#endif