libcxx/include/__random/discrete_distribution.h

   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_DISCRETE_DISTRIBUTION_H
  10 #define _LIBCPP___RANDOM_DISCRETE_DISTRIBUTION_H
  11
  12 #include <__algorithm/upper_bound.h>
  13 #include <__config>
  14 #include <__random/is_valid.h>
  15 #include <__random/uniform_real_distribution.h>
  16 #include <cstddef>
  17 #include <iosfwd>
  18 #include <numeric>
  19 #include <vector>
  20
  21 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
  22 #  pragma GCC system_header
  23 #endif
  24
  25 _LIBCPP_PUSH_MACROS
  26 #include <__undef_macros>
  27
  28 _LIBCPP_BEGIN_NAMESPACE_STD
  29
  30 template<class _IntType = int>
  31 class _LIBCPP_TEMPLATE_VIS discrete_distribution
  32 {
  33     static_assert(__libcpp_random_is_valid_inttype<_IntType>::value, "IntType must be a supported integer type");
  34 public:
  35     // types
  36     typedef _IntType result_type;
  37
  38     class _LIBCPP_TEMPLATE_VIS param_type
  39     {
  40         vector<double> __p_;
  41     public:
  42         typedef discrete_distribution distribution_type;
  43
  44         _LIBCPP_INLINE_VISIBILITY
  45         param_type() {}
  46         template<class _InputIterator>
  47             _LIBCPP_INLINE_VISIBILITY
  48             param_type(_InputIterator __f, _InputIterator __l)
  49             : __p_(__f, __l) {__init();}
  50 #ifndef _LIBCPP_CXX03_LANG
  51         _LIBCPP_INLINE_VISIBILITY
  52         param_type(initializer_list<double> __wl)
  53             : __p_(__wl.begin(), __wl.end()) {__init();}
  54 #endif // _LIBCPP_CXX03_LANG
  55         template<class _UnaryOperation>
  56         _LIBCPP_HIDE_FROM_ABI param_type(size_t __nw, double __xmin, double __xmax,
  57                        _UnaryOperation __fw);
  58
  59         _LIBCPP_HIDE_FROM_ABI vector<double> probabilities() const;
  60
  61         friend _LIBCPP_INLINE_VISIBILITY
  62             bool operator==(const param_type& __x, const param_type& __y)
  63             {return __x.__p_ == __y.__p_;}
  64         friend _LIBCPP_INLINE_VISIBILITY
  65             bool operator!=(const param_type& __x, const param_type& __y)
  66             {return !(__x == __y);}
  67
  68     private:
  69         _LIBCPP_HIDE_FROM_ABI void __init();
  70
  71         friend class discrete_distribution;
  72
  73         template <class _CharT, class _Traits, class _IT>
  74         friend
  75         basic_ostream<_CharT, _Traits>&
  76         operator<<(basic_ostream<_CharT, _Traits>& __os,
  77                    const discrete_distribution<_IT>& __x);
  78
  79         template <class _CharT, class _Traits, class _IT>
  80         friend
  81         basic_istream<_CharT, _Traits>&
  82         operator>>(basic_istream<_CharT, _Traits>& __is,
  83                    discrete_distribution<_IT>& __x);
  84     };
  85
  86 private:
  87     param_type __p_;
  88
  89 public:
  90     // constructor and reset functions
  91     _LIBCPP_INLINE_VISIBILITY
  92     discrete_distribution() {}
  93     template<class _InputIterator>
  94         _LIBCPP_INLINE_VISIBILITY
  95         discrete_distribution(_InputIterator __f, _InputIterator __l)
  96             : __p_(__f, __l) {}
  97 #ifndef _LIBCPP_CXX03_LANG
  98     _LIBCPP_INLINE_VISIBILITY
  99     discrete_distribution(initializer_list<double> __wl)
 100         : __p_(__wl) {}
 101 #endif // _LIBCPP_CXX03_LANG
 102     template<class _UnaryOperation>
 103         _LIBCPP_INLINE_VISIBILITY
 104         discrete_distribution(size_t __nw, double __xmin, double __xmax,
 105                               _UnaryOperation __fw)
 106         : __p_(__nw, __xmin, __xmax, __fw) {}
 107     _LIBCPP_INLINE_VISIBILITY
 108     explicit discrete_distribution(const param_type& __p)
 109         : __p_(__p) {}
 110     _LIBCPP_INLINE_VISIBILITY
 111     void reset() {}
 112
 113     // generating functions
 114     template<class _URNG>
 115         _LIBCPP_INLINE_VISIBILITY
 116         result_type operator()(_URNG& __g)
 117         {return (*this)(__g, __p_);}
 118     template<class _URNG>
 119     _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);
 120
 121     // property functions
 122     _LIBCPP_INLINE_VISIBILITY
 123     vector<double> probabilities() const {return __p_.probabilities();}
 124
 125     _LIBCPP_INLINE_VISIBILITY
 126     param_type param() const {return __p_;}
 127     _LIBCPP_INLINE_VISIBILITY
 128     void param(const param_type& __p) {__p_ = __p;}
 129
 130     _LIBCPP_INLINE_VISIBILITY
 131     result_type min() const {return 0;}
 132     _LIBCPP_INLINE_VISIBILITY
 133     result_type max() const {return __p_.__p_.size();}
 134
 135     friend _LIBCPP_INLINE_VISIBILITY
 136         bool operator==(const discrete_distribution& __x,
 137                         const discrete_distribution& __y)
 138         {return __x.__p_ == __y.__p_;}
 139     friend _LIBCPP_INLINE_VISIBILITY
 140         bool operator!=(const discrete_distribution& __x,
 141                         const discrete_distribution& __y)
 142         {return !(__x == __y);}
 143
 144     template <class _CharT, class _Traits, class _IT>
 145     friend
 146     basic_ostream<_CharT, _Traits>&
 147     operator<<(basic_ostream<_CharT, _Traits>& __os,
 148                const discrete_distribution<_IT>& __x);
 149
 150     template <class _CharT, class _Traits, class _IT>
 151     friend
 152     basic_istream<_CharT, _Traits>&
 153     operator>>(basic_istream<_CharT, _Traits>& __is,
 154                discrete_distribution<_IT>& __x);
 155 };
 156
 157 template<class _IntType>
 158 template<class _UnaryOperation>
 159 discrete_distribution<_IntType>::param_type::param_type(size_t __nw,
 160                                                         double __xmin,
 161                                                         double __xmax,
 162                                                         _UnaryOperation __fw)
 163 {
 164     if (__nw > 1)
 165     {
 166         __p_.reserve(__nw - 1);
 167         double __d = (__xmax - __xmin) / __nw;
 168         double __d2 = __d / 2;
 169         for (size_t __k = 0; __k < __nw; ++__k)
 170             __p_.push_back(__fw(__xmin + __k * __d + __d2));
 171         __init();
 172     }
 173 }
 174
 175 template<class _IntType>
 176 void
 177 discrete_distribution<_IntType>::param_type::__init()
 178 {
 179     if (!__p_.empty())
 180     {
 181         if (__p_.size() > 1)
 182         {
 183             double __s = _VSTD::accumulate(__p_.begin(), __p_.end(), 0.0);
 184             for (vector<double>::iterator __i = __p_.begin(), __e = __p_.end(); __i < __e; ++__i)
 185                 *__i /= __s;
 186             vector<double> __t(__p_.size() - 1);
 187             _VSTD::partial_sum(__p_.begin(), __p_.end() - 1, __t.begin());
 188             swap(__p_, __t);
 189         }
 190         else
 191         {
 192             __p_.clear();
 193             __p_.shrink_to_fit();
 194         }
 195     }
 196 }
 197
 198 template<class _IntType>
 199 vector<double>
 200 discrete_distribution<_IntType>::param_type::probabilities() const
 201 {
 202     size_t __n = __p_.size();
 203     vector<double> __p(__n+1);
 204     _VSTD::adjacent_difference(__p_.begin(), __p_.end(), __p.begin());
 205     if (__n > 0)
 206         __p[__n] = 1 - __p_[__n-1];
 207     else
 208         __p[0] = 1;
 209     return __p;
 210 }
 211
 212 template<class _IntType>
 213 template<class _URNG>
 214 _IntType
 215 discrete_distribution<_IntType>::operator()(_URNG& __g, const param_type& __p)
 216 {
 217     static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
 218     uniform_real_distribution<double> __gen;
 219     return static_cast<_IntType>(
 220            _VSTD::upper_bound(__p.__p_.begin(), __p.__p_.end(), __gen(__g)) -
 221                                                               __p.__p_.begin());
 222 }
 223
 224 template <class _CharT, class _Traits, class _IT>
 225 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
 226 operator<<(basic_ostream<_CharT, _Traits>& __os,
 227            const discrete_distribution<_IT>& __x)
 228 {
 229     __save_flags<_CharT, _Traits> __lx(__os);
 230     typedef basic_ostream<_CharT, _Traits> _OStream;
 231     __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
 232                _OStream::scientific);
 233     _CharT __sp = __os.widen(' ');
 234     __os.fill(__sp);
 235     size_t __n = __x.__p_.__p_.size();
 236     __os << __n;
 237     for (size_t __i = 0; __i < __n; ++__i)
 238         __os << __sp << __x.__p_.__p_[__i];
 239     return __os;
 240 }
 241
 242 template <class _CharT, class _Traits, class _IT>
 243 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
 244 operator>>(basic_istream<_CharT, _Traits>& __is,
 245            discrete_distribution<_IT>& __x)
 246 {
 247     __save_flags<_CharT, _Traits> __lx(__is);
 248     typedef basic_istream<_CharT, _Traits> _Istream;
 249     __is.flags(_Istream::dec | _Istream::skipws);
 250     size_t __n;
 251     __is >> __n;
 252     vector<double> __p(__n);
 253     for (size_t __i = 0; __i < __n; ++__i)
 254         __is >> __p[__i];
 255     if (!__is.fail())
 256         swap(__x.__p_.__p_, __p);
 257     return __is;
 258 }
 259
 260 _LIBCPP_END_NAMESPACE_STD
 261
 262 _LIBCPP_POP_MACROS
 263
 264 #endif // _LIBCPP___RANDOM_DISCRETE_DISTRIBUTION_H