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