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