158 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H10#define _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H11 12#include <__assert>13#include <__config>14#include <__random/bernoulli_distribution.h>15#include <__random/gamma_distribution.h>16#include <__random/is_valid.h>17#include <__random/poisson_distribution.h>18#include <iosfwd>19#include <limits>20 21#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)22# pragma GCC system_header23#endif24 25_LIBCPP_PUSH_MACROS26#include <__undef_macros>27 28_LIBCPP_BEGIN_NAMESPACE_STD29 30template <class _IntType = int>31class negative_binomial_distribution {32 static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");33 34public:35 // types36 typedef _IntType result_type;37 38 class param_type {39 result_type __k_;40 double __p_;41 42 public:43 typedef negative_binomial_distribution distribution_type;44 45 _LIBCPP_HIDE_FROM_ABI explicit param_type(result_type __k = 1, double __p = 0.5) : __k_(__k), __p_(__p) {}46 47 _LIBCPP_HIDE_FROM_ABI result_type k() const { return __k_; }48 _LIBCPP_HIDE_FROM_ABI double p() const { return __p_; }49 50 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {51 return __x.__k_ == __y.__k_ && __x.__p_ == __y.__p_;52 }53 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }54 };55 56private:57 param_type __p_;58 59public:60 // constructor and reset functions61#ifndef _LIBCPP_CXX03_LANG62 _LIBCPP_HIDE_FROM_ABI negative_binomial_distribution() : negative_binomial_distribution(1) {}63 _LIBCPP_HIDE_FROM_ABI explicit negative_binomial_distribution(result_type __k, double __p = 0.5) : __p_(__k, __p) {}64#else65 _LIBCPP_HIDE_FROM_ABI explicit negative_binomial_distribution(result_type __k = 1, double __p = 0.5)66 : __p_(__k, __p) {}67#endif68 _LIBCPP_HIDE_FROM_ABI explicit negative_binomial_distribution(const param_type& __p) : __p_(__p) {}69 _LIBCPP_HIDE_FROM_ABI void reset() {}70 71 // generating functions72 template <class _URNG>73 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {74 return (*this)(__g, __p_);75 }76 template <class _URNG>77 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);78 79 // property functions80 _LIBCPP_HIDE_FROM_ABI result_type k() const { return __p_.k(); }81 _LIBCPP_HIDE_FROM_ABI double p() const { return __p_.p(); }82 83 _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }84 _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }85 86 _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }87 _LIBCPP_HIDE_FROM_ABI result_type max() const { return numeric_limits<result_type>::max(); }88 89 friend _LIBCPP_HIDE_FROM_ABI bool90 operator==(const negative_binomial_distribution& __x, const negative_binomial_distribution& __y) {91 return __x.__p_ == __y.__p_;92 }93 friend _LIBCPP_HIDE_FROM_ABI bool94 operator!=(const negative_binomial_distribution& __x, const negative_binomial_distribution& __y) {95 return !(__x == __y);96 }97};98 99template <class _IntType>100template <class _URNG>101_IntType negative_binomial_distribution<_IntType>::operator()(_URNG& __urng, const param_type& __pr) {102 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");103 result_type __k = __pr.k();104 double __p = __pr.p();105 // When the number of bits in _IntType is small, we are too likely to106 // overflow __f below to use this technique.107 if (__k <= 21 * __p && sizeof(_IntType) > 1) {108 bernoulli_distribution __gen(__p);109 result_type __f = 0;110 result_type __s = 0;111 while (__s < __k) {112 if (__gen(__urng))113 ++__s;114 else115 ++__f;116 }117 _LIBCPP_ASSERT_INTERNAL(__f >= 0,118 "std::negative_binomial_distribution should never produce negative values. "119 "This is almost certainly a signed integer overflow issue on __f.");120 return __f;121 }122 return poisson_distribution<result_type>(gamma_distribution<double>(__k, (1 - __p) / __p)(__urng))(__urng);123}124 125template <class _CharT, class _Traits, class _IntType>126_LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&127operator<<(basic_ostream<_CharT, _Traits>& __os, const negative_binomial_distribution<_IntType>& __x) {128 __save_flags<_CharT, _Traits> __lx(__os);129 typedef basic_ostream<_CharT, _Traits> _OStream;130 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);131 _CharT __sp = __os.widen(' ');132 __os.fill(__sp);133 return __os << __x.k() << __sp << __x.p();134}135 136template <class _CharT, class _Traits, class _IntType>137_LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&138operator>>(basic_istream<_CharT, _Traits>& __is, negative_binomial_distribution<_IntType>& __x) {139 typedef negative_binomial_distribution<_IntType> _Eng;140 typedef typename _Eng::result_type result_type;141 typedef typename _Eng::param_type param_type;142 __save_flags<_CharT, _Traits> __lx(__is);143 typedef basic_istream<_CharT, _Traits> _Istream;144 __is.flags(_Istream::dec | _Istream::skipws);145 result_type __k;146 double __p;147 __is >> __k >> __p;148 if (!__is.fail())149 __x.param(param_type(__k, __p));150 return __is;151}152 153_LIBCPP_END_NAMESPACE_STD154 155_LIBCPP_POP_MACROS156 157#endif // _LIBCPP___RANDOM_NEGATIVE_BINOMIAL_DISTRIBUTION_H158