xref: /freebsd/contrib/llvm-project/libcxx/include/__random/binomial_distribution.h (revision 5956d97f4b3204318ceb6aa9c77bd0bc6ea87a41)
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_BINOMIAL_DISTRIBUTION_H
10 #define _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H
11 
12 #include <__config>
13 #include <__random/uniform_real_distribution.h>
14 #include <cmath>
15 #include <iosfwd>
16 
17 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
18 #pragma GCC system_header
19 #endif
20 
21 _LIBCPP_PUSH_MACROS
22 #include <__undef_macros>
23 
24 _LIBCPP_BEGIN_NAMESPACE_STD
25 
26 template<class _IntType = int>
27 class _LIBCPP_TEMPLATE_VIS binomial_distribution
28 {
29 public:
30     // types
31     typedef _IntType result_type;
32 
33     class _LIBCPP_TEMPLATE_VIS param_type
34     {
35         result_type __t_;
36         double __p_;
37         double __pr_;
38         double __odds_ratio_;
39         result_type __r0_;
40     public:
41         typedef binomial_distribution distribution_type;
42 
43         explicit param_type(result_type __t = 1, double __p = 0.5);
44 
45         _LIBCPP_INLINE_VISIBILITY
46         result_type t() const {return __t_;}
47         _LIBCPP_INLINE_VISIBILITY
48         double p() const {return __p_;}
49 
50         friend _LIBCPP_INLINE_VISIBILITY
51             bool operator==(const param_type& __x, const param_type& __y)
52             {return __x.__t_ == __y.__t_ && __x.__p_ == __y.__p_;}
53         friend _LIBCPP_INLINE_VISIBILITY
54             bool operator!=(const param_type& __x, const param_type& __y)
55             {return !(__x == __y);}
56 
57         friend class binomial_distribution;
58     };
59 
60 private:
61     param_type __p_;
62 
63 public:
64     // constructors and reset functions
65 #ifndef _LIBCPP_CXX03_LANG
66     _LIBCPP_INLINE_VISIBILITY
67     binomial_distribution() : binomial_distribution(1) {}
68     _LIBCPP_INLINE_VISIBILITY
69     explicit binomial_distribution(result_type __t, double __p = 0.5)
70         : __p_(param_type(__t, __p)) {}
71 #else
72     _LIBCPP_INLINE_VISIBILITY
73     explicit binomial_distribution(result_type __t = 1, double __p = 0.5)
74         : __p_(param_type(__t, __p)) {}
75 #endif
76     _LIBCPP_INLINE_VISIBILITY
77     explicit binomial_distribution(const param_type& __p) : __p_(__p) {}
78     _LIBCPP_INLINE_VISIBILITY
79     void reset() {}
80 
81     // generating functions
82     template<class _URNG>
83         _LIBCPP_INLINE_VISIBILITY
84         result_type operator()(_URNG& __g)
85         {return (*this)(__g, __p_);}
86     template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p);
87 
88     // property functions
89     _LIBCPP_INLINE_VISIBILITY
90     result_type t() const {return __p_.t();}
91     _LIBCPP_INLINE_VISIBILITY
92     double p() const {return __p_.p();}
93 
94     _LIBCPP_INLINE_VISIBILITY
95     param_type param() const {return __p_;}
96     _LIBCPP_INLINE_VISIBILITY
97     void param(const param_type& __p) {__p_ = __p;}
98 
99     _LIBCPP_INLINE_VISIBILITY
100     result_type min() const {return 0;}
101     _LIBCPP_INLINE_VISIBILITY
102     result_type max() const {return t();}
103 
104     friend _LIBCPP_INLINE_VISIBILITY
105         bool operator==(const binomial_distribution& __x,
106                         const binomial_distribution& __y)
107         {return __x.__p_ == __y.__p_;}
108     friend _LIBCPP_INLINE_VISIBILITY
109         bool operator!=(const binomial_distribution& __x,
110                         const binomial_distribution& __y)
111         {return !(__x == __y);}
112 };
113 
114 #ifndef _LIBCPP_MSVCRT_LIKE
115 extern "C" double lgamma_r(double, int *);
116 #endif
117 
118 inline _LIBCPP_INLINE_VISIBILITY double __libcpp_lgamma(double __d) {
119 #if defined(_LIBCPP_MSVCRT_LIKE)
120   return lgamma(__d);
121 #else
122   int __sign;
123   return lgamma_r(__d, &__sign);
124 #endif
125 }
126 
127 template<class _IntType>
128 binomial_distribution<_IntType>::param_type::param_type(result_type __t, double __p)
129     : __t_(__t), __p_(__p)
130 {
131     if (0 < __p_ && __p_ < 1)
132     {
133         __r0_ = static_cast<result_type>((__t_ + 1) * __p_);
134         __pr_ = _VSTD::exp(__libcpp_lgamma(__t_ + 1.) -
135                            __libcpp_lgamma(__r0_ + 1.) -
136                            __libcpp_lgamma(__t_ - __r0_ + 1.) + __r0_ * _VSTD::log(__p_) +
137                            (__t_ - __r0_) * _VSTD::log(1 - __p_));
138         __odds_ratio_ = __p_ / (1 - __p_);
139     }
140 }
141 
142 // Reference: Kemp, C.D. (1986). `A modal method for generating binomial
143 //           variables', Commun. Statist. - Theor. Meth. 15(3), 805-813.
144 template<class _IntType>
145 template<class _URNG>
146 _IntType
147 binomial_distribution<_IntType>::operator()(_URNG& __g, const param_type& __pr)
148 {
149     if (__pr.__t_ == 0 || __pr.__p_ == 0)
150         return 0;
151     if (__pr.__p_ == 1)
152         return __pr.__t_;
153     uniform_real_distribution<double> __gen;
154     double __u = __gen(__g) - __pr.__pr_;
155     if (__u < 0)
156         return __pr.__r0_;
157     double __pu = __pr.__pr_;
158     double __pd = __pu;
159     result_type __ru = __pr.__r0_;
160     result_type __rd = __ru;
161     while (true)
162     {
163         bool __break = true;
164         if (__rd >= 1)
165         {
166             __pd *= __rd / (__pr.__odds_ratio_ * (__pr.__t_ - __rd + 1));
167             __u -= __pd;
168             __break = false;
169             if (__u < 0)
170                 return __rd - 1;
171         }
172         if ( __rd != 0 )
173             --__rd;
174         ++__ru;
175         if (__ru <= __pr.__t_)
176         {
177             __pu *= (__pr.__t_ - __ru + 1) * __pr.__odds_ratio_ / __ru;
178             __u -= __pu;
179             __break = false;
180             if (__u < 0)
181                 return __ru;
182         }
183         if (__break)
184             return 0;
185     }
186 }
187 
188 template <class _CharT, class _Traits, class _IntType>
189 basic_ostream<_CharT, _Traits>&
190 operator<<(basic_ostream<_CharT, _Traits>& __os,
191            const binomial_distribution<_IntType>& __x)
192 {
193     __save_flags<_CharT, _Traits> __lx(__os);
194     typedef basic_ostream<_CharT, _Traits> _OStream;
195     __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
196                _OStream::scientific);
197     _CharT __sp = __os.widen(' ');
198     __os.fill(__sp);
199     return __os << __x.t() << __sp << __x.p();
200 }
201 
202 template <class _CharT, class _Traits, class _IntType>
203 basic_istream<_CharT, _Traits>&
204 operator>>(basic_istream<_CharT, _Traits>& __is,
205            binomial_distribution<_IntType>& __x)
206 {
207     typedef binomial_distribution<_IntType> _Eng;
208     typedef typename _Eng::result_type result_type;
209     typedef typename _Eng::param_type param_type;
210     __save_flags<_CharT, _Traits> __lx(__is);
211     typedef basic_istream<_CharT, _Traits> _Istream;
212     __is.flags(_Istream::dec | _Istream::skipws);
213     result_type __t;
214     double __p;
215     __is >> __t >> __p;
216     if (!__is.fail())
217         __x.param(param_type(__t, __p));
218     return __is;
219 }
220 
221 _LIBCPP_END_NAMESPACE_STD
222 
223 _LIBCPP_POP_MACROS
224 
225 #endif // _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H
226