<commit_before>#pragma once
#ifndef WTL_RANDOM_HPP_
#define WTL_RANDOM_HPP_

#include <cstdint>
#include <limits>
#include <random>
#include <vector>
#include <unordered_set>

/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////
namespace wtl {
/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////

namespace {
    union bits64_t {
        uint64_t as_uint64_t;
        uint32_t as_uint32_t[2];
        double as_double;

        bits64_t(uint64_t x): as_uint64_t{x} {}
        bits64_t(uint32_t x, uint32_t y): as_uint32_t{x, y} {}

        // Use 52 bits to make double [0.0, 1.0)
        double as_canonical() {
            as_uint64_t >>= 2;
            as_uint64_t |= 0x3ff0'0000'0000'0000u;
            return as_double -= 1.0;
        }
    };
}

class random_device_64 {
  public:
    using result_type = uint64_t;
    static constexpr result_type min() {return 0u;}
    static constexpr result_type max() {return std::numeric_limits<result_type>::max();}

    result_type operator()() {
        return bits64_t(std_rd(), std_rd()).as_uint64_t;
    }
    double entropy() const noexcept {return std_rd.entropy();}
  private:
    std::random_device std_rd;
};

template <class URBG> inline
double generate_canonical(URBG& gen) {
    if (URBG::max() == std::numeric_limits<uint64_t>::max()) {
        return bits64_t(gen()).as_canonical();
    } else if (URBG::max() == std::numeric_limits<uint32_t>::max()) {
        return bits64_t(gen(), gen()).as_canonical();
    } else {
        return std::generate_canonical<double, std::numeric_limits<double>::digits>(gen);
    }
}

template <class Iter, class URBG> inline
Iter choice(Iter begin_, Iter end_, URBG& engine) {
    using diff_t = decltype(std::distance(begin_, end_));
    std::uniform_int_distribution<diff_t> uniform(0, std::distance(begin_, end_) - 1);
    std::advance(begin_, uniform(engine));
    return begin_;
}

//! Floyd's algorithm
//! fast if k << n
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_floyd(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::unordered_set<size_t> existing_indices;
    std::vector<typename Container::value_type> dst;
    dst.reserve(k);
    std::uniform_int_distribution<size_t> uniform(0, n - 1);
    for (size_t upper = n - k; upper < n; ++upper) {
        size_t idx = std::uniform_int_distribution<size_t>(0, upper)(engine);
        if (existing_indices.insert(idx).second) {
            dst.push_back(src[idx]);
        } else {
            dst.push_back(src[upper]);
        }
    }
    return dst;
}

//! Fisher-Yates algorithm; Durstenfeld; Knuth's P
//! consistently fast; note that whole src is copied first
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_fisher(const Container& src, size_t k, URBG& engine) {
    std::vector<typename Container::value_type> dst(std::begin(src), std::end(src));
    const size_t n = dst.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    const size_t last = n - 1;
    for (size_t i=0; i<k; ++i) {
        std::swap(dst[std::uniform_int_distribution<size_t>(i, last)(engine)], dst[i]);
    }
    dst.resize(k);
    return dst;
}

//! Knuth's algorithm S
//! fast if k / n is large
//! The order is not random
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_knuth(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::vector<typename Container::value_type> dst;
    dst.reserve(k);
    size_t i = 0;
    for (; i<k; ++i) {
        dst.push_back(src[i]);
    }
    std::uniform_int_distribution<size_t> uniform(0, k - 1);
    while (i < n) {
        if (std::bernoulli_distribution((1.0 / ++i) * k)(engine)) {
            dst[uniform(engine)] = src[i - 1];
        }
    }
    return dst;
}

//! sample integers from [0, n) without replacement
template <class T, class URBG> inline
std::unordered_set<T> sample(T n, T k, URBG& engine) {
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::unordered_set<T> existing_indices;
    for (T upper = n - k; upper < n; ++upper) {
        T x = std::uniform_int_distribution<T>(0, upper)(engine);
        if (!existing_indices.insert(x).second) {
            existing_indices.insert(upper);
        }
    }
    return existing_indices;
}

template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (100 * k < n) {return sample_floyd(src, k, engine);}
    else if (5 * k < n) {return sample_fisher(src, k, engine);}
    else {return sample_knuth(src, k, engine);}
}

/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////

//! Pythonic engine object
template <class Generator>
class Prandom{
  public:
    // alias
    using result_type = unsigned int;
    using argument_type = result_type;

    static constexpr result_type min() {return Generator::min();}
    static constexpr result_type max() {return Generator::max();}

    // constructors
    explicit Prandom(result_type s=std::random_device{}())
    : seed_(s), generator_(s) {seed(s);}

    Prandom(const Prandom&) = delete;<commit_msg>:sparkles: Add negative binomial parameterized with k and mu<commit_after>#pragma once
#ifndef WTL_RANDOM_HPP_
#define WTL_RANDOM_HPP_

#include <cstdint>
#include <limits>
#include <random>
#include <vector>
#include <unordered_set>

/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////
namespace wtl {
/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////

namespace {
    union bits64_t {
        uint64_t as_uint64_t;
        uint32_t as_uint32_t[2];
        double as_double;

        bits64_t(uint64_t x): as_uint64_t{x} {}
        bits64_t(uint32_t x, uint32_t y): as_uint32_t{x, y} {}

        // Use 52 bits to make double [0.0, 1.0)
        double as_canonical() {
            as_uint64_t >>= 2;
            as_uint64_t |= 0x3ff0'0000'0000'0000u;
            return as_double -= 1.0;
        }
    };
}

class random_device_64 {
  public:
    using result_type = uint64_t;
    static constexpr result_type min() {return 0u;}
    static constexpr result_type max() {return std::numeric_limits<result_type>::max();}

    result_type operator()() {
        return bits64_t(std_rd(), std_rd()).as_uint64_t;
    }
    double entropy() const noexcept {return std_rd.entropy();}
  private:
    std::random_device std_rd;
};

template <class URBG> inline
double generate_canonical(URBG& gen) {
    if (URBG::max() == std::numeric_limits<uint64_t>::max()) {
        return bits64_t(gen()).as_canonical();
    } else if (URBG::max() == std::numeric_limits<uint32_t>::max()) {
        return bits64_t(gen(), gen()).as_canonical();
    } else {
        return std::generate_canonical<double, std::numeric_limits<double>::digits>(gen);
    }
}

template <class Iter, class URBG> inline
Iter choice(Iter begin_, Iter end_, URBG& engine) {
    using diff_t = decltype(std::distance(begin_, end_));
    std::uniform_int_distribution<diff_t> uniform(0, std::distance(begin_, end_) - 1);
    std::advance(begin_, uniform(engine));
    return begin_;
}

//! Floyd's algorithm
//! fast if k << n
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_floyd(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::unordered_set<size_t> existing_indices;
    std::vector<typename Container::value_type> dst;
    dst.reserve(k);
    std::uniform_int_distribution<size_t> uniform(0, n - 1);
    for (size_t upper = n - k; upper < n; ++upper) {
        size_t idx = std::uniform_int_distribution<size_t>(0, upper)(engine);
        if (existing_indices.insert(idx).second) {
            dst.push_back(src[idx]);
        } else {
            dst.push_back(src[upper]);
        }
    }
    return dst;
}

//! Fisher-Yates algorithm; Durstenfeld; Knuth's P
//! consistently fast; note that whole src is copied first
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_fisher(const Container& src, size_t k, URBG& engine) {
    std::vector<typename Container::value_type> dst(std::begin(src), std::end(src));
    const size_t n = dst.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    const size_t last = n - 1;
    for (size_t i=0; i<k; ++i) {
        std::swap(dst[std::uniform_int_distribution<size_t>(i, last)(engine)], dst[i]);
    }
    dst.resize(k);
    return dst;
}

//! Knuth's algorithm S
//! fast if k / n is large
//! The order is not random
template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample_knuth(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::vector<typename Container::value_type> dst;
    dst.reserve(k);
    size_t i = 0;
    for (; i<k; ++i) {
        dst.push_back(src[i]);
    }
    std::uniform_int_distribution<size_t> uniform(0, k - 1);
    while (i < n) {
        if (std::bernoulli_distribution((1.0 / ++i) * k)(engine)) {
            dst[uniform(engine)] = src[i - 1];
        }
    }
    return dst;
}

//! sample integers from [0, n) without replacement
template <class T, class URBG> inline
std::unordered_set<T> sample(T n, T k, URBG& engine) {
    if (n < k) throw std::runtime_error(std::string(__PRETTY_FUNCTION__) + ": n < k");
    std::unordered_set<T> existing_indices;
    for (T upper = n - k; upper < n; ++upper) {
        T x = std::uniform_int_distribution<T>(0, upper)(engine);
        if (!existing_indices.insert(x).second) {
            existing_indices.insert(upper);
        }
    }
    return existing_indices;
}

template <class Container, class URBG> inline
std::vector<typename Container::value_type>
sample(const Container& src, size_t k, URBG& engine) {
    const size_t n = src.size();
    if (100 * k < n) {return sample_floyd(src, k, engine);}
    else if (5 * k < n) {return sample_fisher(src, k, engine);}
    else {return sample_knuth(src, k, engine);}
}

// std::geometric_distribution if (k == 1)
// std::poisson_distribution   if (k == \infty)
template <class IntType> inline
std::negative_binomial_distribution<IntType>
make_negative_binomial_distribution(IntType k, double mu) {
    double prob = k / (mu + k);
    return std::negative_binomial_distribution<IntType>(k, prob);
}

/////////1/////////2/////////3/////////4/////////5/////////6/////////7/////////

//! Pythonic engine object
template <class Generator>
class Prandom{
  public:
    // alias
    using result_type = unsigned int;
    using argument_type = result_type;

    static constexpr result_type min() {return Generator::min();}
    static constexpr result_type max() {return Generator::max();}

    // constructors
    explicit Prandom(result_type s=std::random_device{}())
    : seed_(s), generator_(s) {seed(s);}

    Prandom(const Prandom&) = delete;
{'fraction_non_alphanumeric': 0.1753381893860562, 'fraction_numerical': 0.01916406520985085, 'mean_word_length': 3.0637773079633543, 'pattern_counts': {'":': 8, '<': 108, '<?xml version=': 0, '>': 78, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4383631', 'n_tokens_mistral': 4067, 'n_tokens_neox': 3862, 'n_words': 934}