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)
__k_(__k)46 : __k_(__k), __p_(__p) {}
47
48 _LIBCPP_INLINE_VISIBILITY
k()49 result_type k() const {return __k_;}
50 _LIBCPP_INLINE_VISIBILITY
p()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
negative_binomial_distribution()68 negative_binomial_distribution() : negative_binomial_distribution(1) {}
69 _LIBCPP_INLINE_VISIBILITY
70 explicit negative_binomial_distribution(result_type __k, double __p = 0.5)
__p_(__k,__p)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
negative_binomial_distribution(const param_type & __p)79 explicit negative_binomial_distribution(const param_type& __p) : __p_(__p) {}
80 _LIBCPP_INLINE_VISIBILITY
reset()81 void reset() {}
82
83 // generating functions
84 template<class _URNG>
85 _LIBCPP_INLINE_VISIBILITY
operator()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
k()92 result_type k() const {return __p_.k();}
93 _LIBCPP_INLINE_VISIBILITY
p()94 double p() const {return __p_.p();}
95
96 _LIBCPP_INLINE_VISIBILITY
param()97 param_type param() const {return __p_;}
98 _LIBCPP_INLINE_VISIBILITY
param(const param_type & __p)99 void param(const param_type& __p) {__p_ = __p;}
100
101 _LIBCPP_INLINE_VISIBILITY
min()102 result_type min() const {return 0;}
103 _LIBCPP_INLINE_VISIBILITY
max()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
operator()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