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 <__math/exponential_functions.h>
14#include <__math/gamma.h>
15#include <__math/logarithms.h>
16#include <__random/is_valid.h>
17#include <__random/uniform_real_distribution.h>
18#include <iosfwd>
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
29template <class _IntType = int>
30class binomial_distribution {
31 static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");
32
33public:
34 // types
35 typedef _IntType result_type;
36
37 class param_type {
38 result_type __t_;
39 double __p_;
40 double __pr_;
41 double __odds_ratio_;
42 result_type __r0_;
43
44 public:
45 typedef binomial_distribution distribution_type;
46
47 _LIBCPP_HIDE_FROM_ABI explicit param_type(result_type __t = 1, double __p = 0.5);
48
49 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type t() const { return __t_; }
50 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI double p() const { return __p_; }
51
52 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {
53 return __x.__t_ == __y.__t_ && __x.__p_ == __y.__p_;
54 }
55 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }
56
57 friend class binomial_distribution;
58 };
59
60private:
61 param_type __p_;
62
63public:
64 // constructors and reset functions
65#ifndef _LIBCPP_CXX03_LANG
66 _LIBCPP_HIDE_FROM_ABI binomial_distribution() : binomial_distribution(1) {}
67 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t, double __p = 0.5)
68 : __p_(param_type(__t, __p)) {}
69#else
70 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(result_type __t = 1, double __p = 0.5)
71 : __p_(param_type(__t, __p)) {}
72#endif
73 _LIBCPP_HIDE_FROM_ABI explicit binomial_distribution(const param_type& __p) : __p_(__p) {}
74 _LIBCPP_HIDE_FROM_ABI void reset() {}
75
76 // generating functions
77 template <class _URNG>
78 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {
79 return (*this)(__g, __p_);
80 }
81 template <class _URNG>
82 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);
83
84 // property functions
85 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type t() const { return __p_.t(); }
86 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI double p() const { return __p_.p(); }
87
88 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }
89 _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }
90
91 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }
92 [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type max() const { return t(); }
93
94 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const binomial_distribution& __x, const binomial_distribution& __y) {
95 return __x.__p_ == __y.__p_;
96 }
97 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const binomial_distribution& __x, const binomial_distribution& __y) {
98 return !(__x == __y);
99 }
100};
101
102template <class _IntType>
103binomial_distribution<_IntType>::param_type::param_type(result_type __t, double __p) : __t_(__t), __p_(__p) {
104 if (0 < __p_ && __p_ < 1) {
105 __r0_ = static_cast<result_type>((__t_ + 1) * __p_);
106 __pr_ = __math::exp(
107 __math::__lgamma_r(d: __t_ + 1.) - __math::__lgamma_r(d: __r0_ + 1.) - __math::__lgamma_r(d: __t_ - __r0_ + 1.) +
108 __r0_ * __math::log(x: __p_) + (__t_ - __r0_) * __math::log(x: 1 - __p_));
109 __odds_ratio_ = __p_ / (1 - __p_);
110 }
111}
112
113// Reference: Kemp, C.D. (1986). `A modal method for generating binomial
114// variables', Commun. Statist. - Theor. Meth. 15(3), 805-813.
115template <class _IntType>
116template <class _URNG>
117_IntType binomial_distribution<_IntType>::operator()(_URNG& __g, const param_type& __pr) {
118 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
119 if (__pr.__t_ == 0 || __pr.__p_ == 0)
120 return 0;
121 if (__pr.__p_ == 1)
122 return __pr.__t_;
123 uniform_real_distribution<double> __gen;
124 double __u = __gen(__g) - __pr.__pr_;
125 if (__u < 0)
126 return __pr.__r0_;
127 double __pu = __pr.__pr_;
128 double __pd = __pu;
129 result_type __ru = __pr.__r0_;
130 result_type __rd = __ru;
131 while (true) {
132 bool __break = true;
133 if (__rd >= 1) {
134 __pd *= __rd / (__pr.__odds_ratio_ * (__pr.__t_ - __rd + 1));
135 __u -= __pd;
136 __break = false;
137 if (__u < 0)
138 return __rd - 1;
139 }
140 if (__rd != 0)
141 --__rd;
142 ++__ru;
143 if (__ru <= __pr.__t_) {
144 __pu *= (__pr.__t_ - __ru + 1) * __pr.__odds_ratio_ / __ru;
145 __u -= __pu;
146 __break = false;
147 if (__u < 0)
148 return __ru;
149 }
150 if (__break)
151 return 0;
152 }
153}
154
155template <class _CharT, class _Traits, class _IntType>
156_LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
157operator<<(basic_ostream<_CharT, _Traits>& __os, const binomial_distribution<_IntType>& __x) {
158 __save_flags<_CharT, _Traits> __lx(__os);
159 typedef basic_ostream<_CharT, _Traits> _OStream;
160 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);
161 _CharT __sp = __os.widen(' ');
162 __os.fill(__sp);
163 return __os << __x.t() << __sp << __x.p();
164}
165
166template <class _CharT, class _Traits, class _IntType>
167_LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
168operator>>(basic_istream<_CharT, _Traits>& __is, binomial_distribution<_IntType>& __x) {
169 typedef binomial_distribution<_IntType> _Eng;
170 typedef typename _Eng::result_type result_type;
171 typedef typename _Eng::param_type param_type;
172 __save_flags<_CharT, _Traits> __lx(__is);
173 typedef basic_istream<_CharT, _Traits> _Istream;
174 __is.flags(_Istream::dec | _Istream::skipws);
175 result_type __t;
176 double __p;
177 __is >> __t >> __p;
178 if (!__is.fail())
179 __x.param(param_type(__t, __p));
180 return __is;
181}
182
183_LIBCPP_END_NAMESPACE_STD
184
185_LIBCPP_POP_MACROS
186
187#endif // _LIBCPP___RANDOM_BINOMIAL_DISTRIBUTION_H
188