geometric_distribution.h 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  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. #ifndef _LIBCPP___RANDOM_GEOMETRIC_DISTRIBUTION_H
  9. #define _LIBCPP___RANDOM_GEOMETRIC_DISTRIBUTION_H
  10. #include <__config>
  11. #include <__random/is_valid.h>
  12. #include <__random/negative_binomial_distribution.h>
  13. #include <iosfwd>
  14. #include <limits>
  15. #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
  16. # pragma GCC system_header
  17. #endif
  18. _LIBCPP_PUSH_MACROS
  19. #include <__undef_macros>
  20. _LIBCPP_BEGIN_NAMESPACE_STD
  21. template <class _IntType = int>
  22. class _LIBCPP_TEMPLATE_VIS geometric_distribution {
  23. static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");
  24. public:
  25. // types
  26. typedef _IntType result_type;
  27. class _LIBCPP_TEMPLATE_VIS param_type {
  28. double __p_;
  29. public:
  30. typedef geometric_distribution distribution_type;
  31. _LIBCPP_HIDE_FROM_ABI explicit param_type(double __p = 0.5) : __p_(__p) {}
  32. _LIBCPP_HIDE_FROM_ABI double p() const { return __p_; }
  33. friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {
  34. return __x.__p_ == __y.__p_;
  35. }
  36. friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }
  37. };
  38. private:
  39. param_type __p_;
  40. public:
  41. // constructors and reset functions
  42. #ifndef _LIBCPP_CXX03_LANG
  43. _LIBCPP_HIDE_FROM_ABI geometric_distribution() : geometric_distribution(0.5) {}
  44. _LIBCPP_HIDE_FROM_ABI explicit geometric_distribution(double __p) : __p_(__p) {}
  45. #else
  46. _LIBCPP_HIDE_FROM_ABI explicit geometric_distribution(double __p = 0.5) : __p_(__p) {}
  47. #endif
  48. _LIBCPP_HIDE_FROM_ABI explicit geometric_distribution(const param_type& __p) : __p_(__p) {}
  49. _LIBCPP_HIDE_FROM_ABI void reset() {}
  50. // generating functions
  51. template <class _URNG>
  52. _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {
  53. return (*this)(__g, __p_);
  54. }
  55. template <class _URNG>
  56. _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p) {
  57. return negative_binomial_distribution<result_type>(1, __p.p())(__g);
  58. }
  59. // property functions
  60. _LIBCPP_HIDE_FROM_ABI double p() const { return __p_.p(); }
  61. _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }
  62. _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }
  63. _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }
  64. _LIBCPP_HIDE_FROM_ABI result_type max() const { return numeric_limits<result_type>::max(); }
  65. friend _LIBCPP_HIDE_FROM_ABI bool operator==(const geometric_distribution& __x, const geometric_distribution& __y) {
  66. return __x.__p_ == __y.__p_;
  67. }
  68. friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const geometric_distribution& __x, const geometric_distribution& __y) {
  69. return !(__x == __y);
  70. }
  71. };
  72. template <class _CharT, class _Traits, class _IntType>
  73. _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
  74. operator<<(basic_ostream<_CharT, _Traits>& __os, const geometric_distribution<_IntType>& __x) {
  75. __save_flags<_CharT, _Traits> __lx(__os);
  76. typedef basic_ostream<_CharT, _Traits> _OStream;
  77. __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);
  78. return __os << __x.p();
  79. }
  80. template <class _CharT, class _Traits, class _IntType>
  81. _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
  82. operator>>(basic_istream<_CharT, _Traits>& __is, geometric_distribution<_IntType>& __x) {
  83. typedef geometric_distribution<_IntType> _Eng;
  84. typedef typename _Eng::param_type param_type;
  85. __save_flags<_CharT, _Traits> __lx(__is);
  86. typedef basic_istream<_CharT, _Traits> _Istream;
  87. __is.flags(_Istream::dec | _Istream::skipws);
  88. double __p;
  89. __is >> __p;
  90. if (!__is.fail())
  91. __x.param(param_type(__p));
  92. return __is;
  93. }
  94. _LIBCPP_END_NAMESPACE_STD
  95. _LIBCPP_POP_MACROS
  96. #endif // _LIBCPP___RANDOM_GEOMETRIC_DISTRIBUTION_H