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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #ifndef _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H 10 #define _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H 11 12 #include <__config> 13 #include <__random/bernoulli_distribution.h> 14 #include <__random/gamma_distribution.h> 15 #include <__random/is_valid.h> 16 #include <__random/poisson_distribution.h> 17 #include <iosfwd> 18 #include <limits> 19 20 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER) 21 # pragma GCC system_header 22 #endif 23 24 _LIBCPP_PUSH_MACROS 25 #include <__undef_macros> 26 27 _LIBCPP_BEGIN_NAMESPACE_STD 28 29 template<class _IntType = int> 30 class _LIBCPP_TEMPLATE_VIS negative_binomial_distribution 31 { 32 static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type"); 33 public: 34 // types 35 typedef _IntType result_type; 36 37 class _LIBCPP_TEMPLATE_VIS param_type 38 { 39 result_type __k_; 40 double __p_; 41 public: 42 typedef negative_binomial_distribution distribution_type; 43 44 _LIBCPP_INLINE_VISIBILITY 45 explicit param_type(result_type __k = 1, double __p = 0.5) 46 : __k_(__k), __p_(__p) {} 47 48 _LIBCPP_INLINE_VISIBILITY 49 result_type k() const {return __k_;} 50 _LIBCPP_INLINE_VISIBILITY 51 double p() const {return __p_;} 52 53 friend _LIBCPP_INLINE_VISIBILITY 54 bool operator==(const param_type& __x, const param_type& __y) 55 {return __x.__k_ == __y.__k_ && __x.__p_ == __y.__p_;} 56 friend _LIBCPP_INLINE_VISIBILITY 57 bool operator!=(const param_type& __x, const param_type& __y) 58 {return !(__x == __y);} 59 }; 60 61 private: 62 param_type __p_; 63 64 public: 65 // constructor and reset functions 66 #ifndef _LIBCPP_CXX03_LANG 67 _LIBCPP_INLINE_VISIBILITY 68 negative_binomial_distribution() : negative_binomial_distribution(1) {} 69 _LIBCPP_INLINE_VISIBILITY 70 explicit negative_binomial_distribution(result_type __k, double __p = 0.5) 71 : __p_(__k, __p) {} 72 #else 73 _LIBCPP_INLINE_VISIBILITY 74 explicit negative_binomial_distribution(result_type __k = 1, 75 double __p = 0.5) 76 : __p_(__k, __p) {} 77 #endif 78 _LIBCPP_INLINE_VISIBILITY 79 explicit negative_binomial_distribution(const param_type& __p) : __p_(__p) {} 80 _LIBCPP_INLINE_VISIBILITY 81 void reset() {} 82 83 // generating functions 84 template<class _URNG> 85 _LIBCPP_INLINE_VISIBILITY 86 result_type operator()(_URNG& __g) 87 {return (*this)(__g, __p_);} 88 template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p); 89 90 // property functions 91 _LIBCPP_INLINE_VISIBILITY 92 result_type k() const {return __p_.k();} 93 _LIBCPP_INLINE_VISIBILITY 94 double p() const {return __p_.p();} 95 96 _LIBCPP_INLINE_VISIBILITY 97 param_type param() const {return __p_;} 98 _LIBCPP_INLINE_VISIBILITY 99 void param(const param_type& __p) {__p_ = __p;} 100 101 _LIBCPP_INLINE_VISIBILITY 102 result_type min() const {return 0;} 103 _LIBCPP_INLINE_VISIBILITY 104 result_type max() const {return numeric_limits<result_type>::max();} 105 106 friend _LIBCPP_INLINE_VISIBILITY 107 bool operator==(const negative_binomial_distribution& __x, 108 const negative_binomial_distribution& __y) 109 {return __x.__p_ == __y.__p_;} 110 friend _LIBCPP_INLINE_VISIBILITY 111 bool operator!=(const negative_binomial_distribution& __x, 112 const negative_binomial_distribution& __y) 113 {return !(__x == __y);} 114 }; 115 116 template <class _IntType> 117 template<class _URNG> 118 _IntType 119 negative_binomial_distribution<_IntType>::operator()(_URNG& __urng, const param_type& __pr) 120 { 121 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, ""); 122 result_type __k = __pr.k(); 123 double __p = __pr.p(); 124 // When the number of bits in _IntType is small, we are too likely to 125 // overflow __f below to use this technique. 126 if (__k <= 21 * __p && sizeof(_IntType) > 1) 127 { 128 bernoulli_distribution __gen(__p); 129 result_type __f = 0; 130 result_type __s = 0; 131 while (__s < __k) 132 { 133 if (__gen(__urng)) 134 ++__s; 135 else 136 ++__f; 137 } 138 _LIBCPP_ASSERT(__f >= 0, "std::negative_binomial_distribution should never produce negative values. " 139 "This is almost certainly a signed integer overflow issue on __f."); 140 return __f; 141 } 142 return poisson_distribution<result_type>(gamma_distribution<double> 143 (__k, (1-__p)/__p)(__urng))(__urng); 144 } 145 146 template <class _CharT, class _Traits, class _IntType> 147 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>& 148 operator<<(basic_ostream<_CharT, _Traits>& __os, 149 const negative_binomial_distribution<_IntType>& __x) 150 { 151 __save_flags<_CharT, _Traits> __lx(__os); 152 typedef basic_ostream<_CharT, _Traits> _OStream; 153 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | 154 _OStream::scientific); 155 _CharT __sp = __os.widen(' '); 156 __os.fill(__sp); 157 return __os << __x.k() << __sp << __x.p(); 158 } 159 160 template <class _CharT, class _Traits, class _IntType> 161 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>& 162 operator>>(basic_istream<_CharT, _Traits>& __is, 163 negative_binomial_distribution<_IntType>& __x) 164 { 165 typedef negative_binomial_distribution<_IntType> _Eng; 166 typedef typename _Eng::result_type result_type; 167 typedef typename _Eng::param_type param_type; 168 __save_flags<_CharT, _Traits> __lx(__is); 169 typedef basic_istream<_CharT, _Traits> _Istream; 170 __is.flags(_Istream::dec | _Istream::skipws); 171 result_type __k; 172 double __p; 173 __is >> __k >> __p; 174 if (!__is.fail()) 175 __x.param(param_type(__k, __p)); 176 return __is; 177 } 178 179 _LIBCPP_END_NAMESPACE_STD 180 181 _LIBCPP_POP_MACROS 182 183 #endif // _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H 184