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> 89 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p); 90 91 // property functions 92 _LIBCPP_INLINE_VISIBILITY 93 result_type k() const {return __p_.k();} 94 _LIBCPP_INLINE_VISIBILITY 95 double p() const {return __p_.p();} 96 97 _LIBCPP_INLINE_VISIBILITY 98 param_type param() const {return __p_;} 99 _LIBCPP_INLINE_VISIBILITY 100 void param(const param_type& __p) {__p_ = __p;} 101 102 _LIBCPP_INLINE_VISIBILITY 103 result_type min() const {return 0;} 104 _LIBCPP_INLINE_VISIBILITY 105 result_type max() const {return numeric_limits<result_type>::max();} 106 107 friend _LIBCPP_INLINE_VISIBILITY 108 bool operator==(const negative_binomial_distribution& __x, 109 const negative_binomial_distribution& __y) 110 {return __x.__p_ == __y.__p_;} 111 friend _LIBCPP_INLINE_VISIBILITY 112 bool operator!=(const negative_binomial_distribution& __x, 113 const negative_binomial_distribution& __y) 114 {return !(__x == __y);} 115 }; 116 117 template <class _IntType> 118 template<class _URNG> 119 _IntType 120 negative_binomial_distribution<_IntType>::operator()(_URNG& __urng, const param_type& __pr) 121 { 122 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, ""); 123 result_type __k = __pr.k(); 124 double __p = __pr.p(); 125 // When the number of bits in _IntType is small, we are too likely to 126 // overflow __f below to use this technique. 127 if (__k <= 21 * __p && sizeof(_IntType) > 1) 128 { 129 bernoulli_distribution __gen(__p); 130 result_type __f = 0; 131 result_type __s = 0; 132 while (__s < __k) 133 { 134 if (__gen(__urng)) 135 ++__s; 136 else 137 ++__f; 138 } 139 _LIBCPP_ASSERT_UNCATEGORIZED(__f >= 0, 140 "std::negative_binomial_distribution should never produce negative values. " 141 "This is almost certainly a signed integer overflow issue on __f."); 142 return __f; 143 } 144 return poisson_distribution<result_type>(gamma_distribution<double> 145 (__k, (1-__p)/__p)(__urng))(__urng); 146 } 147 148 template <class _CharT, class _Traits, class _IntType> 149 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>& 150 operator<<(basic_ostream<_CharT, _Traits>& __os, 151 const negative_binomial_distribution<_IntType>& __x) 152 { 153 __save_flags<_CharT, _Traits> __lx(__os); 154 typedef basic_ostream<_CharT, _Traits> _OStream; 155 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | 156 _OStream::scientific); 157 _CharT __sp = __os.widen(' '); 158 __os.fill(__sp); 159 return __os << __x.k() << __sp << __x.p(); 160 } 161 162 template <class _CharT, class _Traits, class _IntType> 163 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>& 164 operator>>(basic_istream<_CharT, _Traits>& __is, 165 negative_binomial_distribution<_IntType>& __x) 166 { 167 typedef negative_binomial_distribution<_IntType> _Eng; 168 typedef typename _Eng::result_type result_type; 169 typedef typename _Eng::param_type param_type; 170 __save_flags<_CharT, _Traits> __lx(__is); 171 typedef basic_istream<_CharT, _Traits> _Istream; 172 __is.flags(_Istream::dec | _Istream::skipws); 173 result_type __k; 174 double __p; 175 __is >> __k >> __p; 176 if (!__is.fail()) 177 __x.param(param_type(__k, __p)); 178 return __is; 179 } 180 181 _LIBCPP_END_NAMESPACE_STD 182 183 _LIBCPP_POP_MACROS 184 185 #endif // _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H 186