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>
89 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);
90
91 // property functions
92 _LIBCPP_INLINE_VISIBILITY
k()93 result_type k() const {return __p_.k();}
94 _LIBCPP_INLINE_VISIBILITY
p()95 double p() const {return __p_.p();}
96
97 _LIBCPP_INLINE_VISIBILITY
param()98 param_type param() const {return __p_;}
99 _LIBCPP_INLINE_VISIBILITY
param(const param_type & __p)100 void param(const param_type& __p) {__p_ = __p;}
101
102 _LIBCPP_INLINE_VISIBILITY
min()103 result_type min() const {return 0;}
104 _LIBCPP_INLINE_VISIBILITY
max()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
operator()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(__f >= 0, "std::negative_binomial_distribution should never produce negative values. "
140 "This is almost certainly a signed integer overflow issue on __f.");
141 return __f;
142 }
143 return poisson_distribution<result_type>(gamma_distribution<double>
144 (__k, (1-__p)/__p)(__urng))(__urng);
145 }
146
147 template <class _CharT, class _Traits, class _IntType>
148 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
149 operator<<(basic_ostream<_CharT, _Traits>& __os,
150 const negative_binomial_distribution<_IntType>& __x)
151 {
152 __save_flags<_CharT, _Traits> __lx(__os);
153 typedef basic_ostream<_CharT, _Traits> _OStream;
154 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
155 _OStream::scientific);
156 _CharT __sp = __os.widen(' ');
157 __os.fill(__sp);
158 return __os << __x.k() << __sp << __x.p();
159 }
160
161 template <class _CharT, class _Traits, class _IntType>
162 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
163 operator>>(basic_istream<_CharT, _Traits>& __is,
164 negative_binomial_distribution<_IntType>& __x)
165 {
166 typedef negative_binomial_distribution<_IntType> _Eng;
167 typedef typename _Eng::result_type result_type;
168 typedef typename _Eng::param_type param_type;
169 __save_flags<_CharT, _Traits> __lx(__is);
170 typedef basic_istream<_CharT, _Traits> _Istream;
171 __is.flags(_Istream::dec | _Istream::skipws);
172 result_type __k;
173 double __p;
174 __is >> __k >> __p;
175 if (!__is.fail())
176 __x.param(param_type(__k, __p));
177 return __is;
178 }
179
180 _LIBCPP_END_NAMESPACE_STD
181
182 _LIBCPP_POP_MACROS
183
184 #endif // _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H
185