206 lines · c
1//===----------------------------------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H10#define _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H11 12#include <__config>13#include <__random/is_valid.h>14#include <__random/uniform_real_distribution.h>15#include <cmath>16#include <iosfwd>17 18#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)19# pragma GCC system_header20#endif21 22_LIBCPP_PUSH_MACROS23#include <__undef_macros>24 25_LIBCPP_BEGIN_NAMESPACE_STD26 27template <class _IntType = int>28class binomial_distribution {29 static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");30 31public:32 // types33 typedef _IntType result_type;34 35 class param_type {36 result_type __t_;37 double __p_;38 double __pr_;39 double __odds_ratio_;40 result_type __r0_;41 42 public:43 typedef binomial_distribution distribution_type;44 45 _LIBCPP_HIDE_FROM_ABI explicit param_type(result_type __t = 1, double __p = 0.5);46 47 _LIBCPP_HIDE_FROM_ABI result_type t() const { return __t_; }48 _LIBCPP_HIDE_FROM_ABI double p() const { return __p_; }49 50 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {51 return __x.__t_ == __y.__t_ && __x.__p_ == __y.__p_;52 }53 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }54 55 friend class binomial_distribution;56 };57 58private:59 param_type __p_;60 61public:62 // constructors and reset functions63#ifndef _LIBCPP_CXX03_LANG64 _LIBCPP_HIDE_FROM_ABI binomial_distribution() : binomial_distribution(1) {}65 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t, double __p = 0.5)66 : __p_(param_type(__t, __p)) {}67#else68 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t = 1, double __p = 0.5)69 : __p_(param_type(__t, __p)) {}70#endif71 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(const param_type& __p) : __p_(__p) {}72 _LIBCPP_HIDE_FROM_ABI void reset() {}73 74 // generating functions75 template <class _URNG>76 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {77 return (*this)(__g, __p_);78 }79 template <class _URNG>80 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);81 82 // property functions83 _LIBCPP_HIDE_FROM_ABI result_type t() const { return __p_.t(); }84 _LIBCPP_HIDE_FROM_ABI double p() const { return __p_.p(); }85 86 _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }87 _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }88 89 _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }90 _LIBCPP_HIDE_FROM_ABI result_type max() const { return t(); }91 92 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const binomial_distribution& __x, const binomial_distribution& __y) {93 return __x.__p_ == __y.__p_;94 }95 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const binomial_distribution& __x, const binomial_distribution& __y) {96 return !(__x == __y);97 }98};99 100// Some libc declares the math functions to be `noexcept`.101#if _LIBCPP_GLIBC_PREREQ(2, 8) || defined(__LLVM_LIBC__)102# define _LIBCPP_LGAMMA_R_NOEXCEPT _NOEXCEPT103#else104# define _LIBCPP_LGAMMA_R_NOEXCEPT105#endif106 107#if !defined(_LIBCPP_MSVCRT_LIKE)108extern "C" double lgamma_r(double, int*) _LIBCPP_LGAMMA_R_NOEXCEPT;109#endif110 111inline _LIBCPP_HIDE_FROM_ABI double __libcpp_lgamma(double __d) {112#if defined(_LIBCPP_MSVCRT_LIKE)113 return lgamma(__d);114#else115 int __sign;116 return lgamma_r(__d, &__sign);117#endif118}119 120template <class _IntType>121binomial_distribution<_IntType>::param_type::param_type(result_type __t, double __p) : __t_(__t), __p_(__p) {122 if (0 < __p_ && __p_ < 1) {123 __r0_ = static_cast<result_type>((__t_ + 1) * __p_);124 __pr_ = std::exp(125 std::__libcpp_lgamma(__t_ + 1.) - std::__libcpp_lgamma(__r0_ + 1.) - std::__libcpp_lgamma(__t_ - __r0_ + 1.) +126 __r0_ * std::log(__p_) + (__t_ - __r0_) * std::log(1 - __p_));127 __odds_ratio_ = __p_ / (1 - __p_);128 }129}130 131// Reference: Kemp, C.D. (1986). `A modal method for generating binomial132// variables', Commun. Statist. - Theor. Meth. 15(3), 805-813.133template <class _IntType>134template <class _URNG>135_IntType binomial_distribution<_IntType>::operator()(_URNG& __g, const param_type& __pr) {136 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");137 if (__pr.__t_ == 0 || __pr.__p_ == 0)138 return 0;139 if (__pr.__p_ == 1)140 return __pr.__t_;141 uniform_real_distribution<double> __gen;142 double __u = __gen(__g) - __pr.__pr_;143 if (__u < 0)144 return __pr.__r0_;145 double __pu = __pr.__pr_;146 double __pd = __pu;147 result_type __ru = __pr.__r0_;148 result_type __rd = __ru;149 while (true) {150 bool __break = true;151 if (__rd >= 1) {152 __pd *= __rd / (__pr.__odds_ratio_ * (__pr.__t_ - __rd + 1));153 __u -= __pd;154 __break = false;155 if (__u < 0)156 return __rd - 1;157 }158 if (__rd != 0)159 --__rd;160 ++__ru;161 if (__ru <= __pr.__t_) {162 __pu *= (__pr.__t_ - __ru + 1) * __pr.__odds_ratio_ / __ru;163 __u -= __pu;164 __break = false;165 if (__u < 0)166 return __ru;167 }168 if (__break)169 return 0;170 }171}172 173template <class _CharT, class _Traits, class _IntType>174_LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&175operator<<(basic_ostream<_CharT, _Traits>& __os, const binomial_distribution<_IntType>& __x) {176 __save_flags<_CharT, _Traits> __lx(__os);177 typedef basic_ostream<_CharT, _Traits> _OStream;178 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);179 _CharT __sp = __os.widen(' ');180 __os.fill(__sp);181 return __os << __x.t() << __sp << __x.p();182}183 184template <class _CharT, class _Traits, class _IntType>185_LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&186operator>>(basic_istream<_CharT, _Traits>& __is, binomial_distribution<_IntType>& __x) {187 typedef binomial_distribution<_IntType> _Eng;188 typedef typename _Eng::result_type result_type;189 typedef typename _Eng::param_type param_type;190 __save_flags<_CharT, _Traits> __lx(__is);191 typedef basic_istream<_CharT, _Traits> _Istream;192 __is.flags(_Istream::dec | _Istream::skipws);193 result_type __t;194 double __p;195 __is >> __t >> __p;196 if (!__is.fail())197 __x.param(param_type(__t, __p));198 return __is;199}200 201_LIBCPP_END_NAMESPACE_STD202 203_LIBCPP_POP_MACROS204 205#endif // _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H206