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
param_type()45         param_type() {}
46         template<class _InputIterator>
47             _LIBCPP_INLINE_VISIBILITY
param_type(_InputIterator __f,_InputIterator __l)48             param_type(_InputIterator __f, _InputIterator __l)
49             : __p_(__f, __l) {__init();}
50 #ifndef _LIBCPP_CXX03_LANG
51         _LIBCPP_INLINE_VISIBILITY
param_type(initializer_list<double> __wl)52         param_type(initializer_list<double> __wl)
53             : __p_(__wl.begin(), __wl.end()) {__init();}
54 #endif // _LIBCPP_CXX03_LANG
55         template<class _UnaryOperation>
56             param_type(size_t __nw, double __xmin, double __xmax,
57                        _UnaryOperation __fw);
58 
59         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         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
discrete_distribution()92     discrete_distribution() {}
93     template<class _InputIterator>
94         _LIBCPP_INLINE_VISIBILITY
discrete_distribution(_InputIterator __f,_InputIterator __l)95         discrete_distribution(_InputIterator __f, _InputIterator __l)
96             : __p_(__f, __l) {}
97 #ifndef _LIBCPP_CXX03_LANG
98     _LIBCPP_INLINE_VISIBILITY
discrete_distribution(initializer_list<double> __wl)99     discrete_distribution(initializer_list<double> __wl)
100         : __p_(__wl) {}
101 #endif // _LIBCPP_CXX03_LANG
102     template<class _UnaryOperation>
103         _LIBCPP_INLINE_VISIBILITY
discrete_distribution(size_t __nw,double __xmin,double __xmax,_UnaryOperation __fw)104         discrete_distribution(size_t __nw, double __xmin, double __xmax,
105                               _UnaryOperation __fw)
106         : __p_(__nw, __xmin, __xmax, __fw) {}
107     _LIBCPP_INLINE_VISIBILITY
discrete_distribution(const param_type & __p)108     explicit discrete_distribution(const param_type& __p)
109         : __p_(__p) {}
110     _LIBCPP_INLINE_VISIBILITY
reset()111     void reset() {}
112 
113     // generating functions
114     template<class _URNG>
115         _LIBCPP_INLINE_VISIBILITY
operator()116         result_type operator()(_URNG& __g)
117         {return (*this)(__g, __p_);}
118     template<class _URNG> result_type operator()(_URNG& __g, const param_type& __p);
119 
120     // property functions
121     _LIBCPP_INLINE_VISIBILITY
probabilities()122     vector<double> probabilities() const {return __p_.probabilities();}
123 
124     _LIBCPP_INLINE_VISIBILITY
param()125     param_type param() const {return __p_;}
126     _LIBCPP_INLINE_VISIBILITY
param(const param_type & __p)127     void param(const param_type& __p) {__p_ = __p;}
128 
129     _LIBCPP_INLINE_VISIBILITY
min()130     result_type min() const {return 0;}
131     _LIBCPP_INLINE_VISIBILITY
max()132     result_type max() const {return __p_.__p_.size();}
133 
134     friend _LIBCPP_INLINE_VISIBILITY
135         bool operator==(const discrete_distribution& __x,
136                         const discrete_distribution& __y)
137         {return __x.__p_ == __y.__p_;}
138     friend _LIBCPP_INLINE_VISIBILITY
139         bool operator!=(const discrete_distribution& __x,
140                         const discrete_distribution& __y)
141         {return !(__x == __y);}
142 
143     template <class _CharT, class _Traits, class _IT>
144     friend
145     basic_ostream<_CharT, _Traits>&
146     operator<<(basic_ostream<_CharT, _Traits>& __os,
147                const discrete_distribution<_IT>& __x);
148 
149     template <class _CharT, class _Traits, class _IT>
150     friend
151     basic_istream<_CharT, _Traits>&
152     operator>>(basic_istream<_CharT, _Traits>& __is,
153                discrete_distribution<_IT>& __x);
154 };
155 
156 template<class _IntType>
157 template<class _UnaryOperation>
param_type(size_t __nw,double __xmin,double __xmax,_UnaryOperation __fw)158 discrete_distribution<_IntType>::param_type::param_type(size_t __nw,
159                                                         double __xmin,
160                                                         double __xmax,
161                                                         _UnaryOperation __fw)
162 {
163     if (__nw > 1)
164     {
165         __p_.reserve(__nw - 1);
166         double __d = (__xmax - __xmin) / __nw;
167         double __d2 = __d / 2;
168         for (size_t __k = 0; __k < __nw; ++__k)
169             __p_.push_back(__fw(__xmin + __k * __d + __d2));
170         __init();
171     }
172 }
173 
174 template<class _IntType>
175 void
__init()176 discrete_distribution<_IntType>::param_type::__init()
177 {
178     if (!__p_.empty())
179     {
180         if (__p_.size() > 1)
181         {
182             double __s = _VSTD::accumulate(__p_.begin(), __p_.end(), 0.0);
183             for (vector<double>::iterator __i = __p_.begin(), __e = __p_.end(); __i < __e; ++__i)
184                 *__i /= __s;
185             vector<double> __t(__p_.size() - 1);
186             _VSTD::partial_sum(__p_.begin(), __p_.end() - 1, __t.begin());
187             swap(__p_, __t);
188         }
189         else
190         {
191             __p_.clear();
192             __p_.shrink_to_fit();
193         }
194     }
195 }
196 
197 template<class _IntType>
198 vector<double>
probabilities()199 discrete_distribution<_IntType>::param_type::probabilities() const
200 {
201     size_t __n = __p_.size();
202     vector<double> __p(__n+1);
203     _VSTD::adjacent_difference(__p_.begin(), __p_.end(), __p.begin());
204     if (__n > 0)
205         __p[__n] = 1 - __p_[__n-1];
206     else
207         __p[0] = 1;
208     return __p;
209 }
210 
211 template<class _IntType>
212 template<class _URNG>
213 _IntType
operator()214 discrete_distribution<_IntType>::operator()(_URNG& __g, const param_type& __p)
215 {
216     static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");
217     uniform_real_distribution<double> __gen;
218     return static_cast<_IntType>(
219            _VSTD::upper_bound(__p.__p_.begin(), __p.__p_.end(), __gen(__g)) -
220                                                               __p.__p_.begin());
221 }
222 
223 template <class _CharT, class _Traits, class _IT>
224 _LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&
225 operator<<(basic_ostream<_CharT, _Traits>& __os,
226            const discrete_distribution<_IT>& __x)
227 {
228     __save_flags<_CharT, _Traits> __lx(__os);
229     typedef basic_ostream<_CharT, _Traits> _OStream;
230     __os.flags(_OStream::dec | _OStream::left | _OStream::fixed |
231                _OStream::scientific);
232     _CharT __sp = __os.widen(' ');
233     __os.fill(__sp);
234     size_t __n = __x.__p_.__p_.size();
235     __os << __n;
236     for (size_t __i = 0; __i < __n; ++__i)
237         __os << __sp << __x.__p_.__p_[__i];
238     return __os;
239 }
240 
241 template <class _CharT, class _Traits, class _IT>
242 _LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&
243 operator>>(basic_istream<_CharT, _Traits>& __is,
244            discrete_distribution<_IT>& __x)
245 {
246     __save_flags<_CharT, _Traits> __lx(__is);
247     typedef basic_istream<_CharT, _Traits> _Istream;
248     __is.flags(_Istream::dec | _Istream::skipws);
249     size_t __n;
250     __is >> __n;
251     vector<double> __p(__n);
252     for (size_t __i = 0; __i < __n; ++__i)
253         __is >> __p[__i];
254     if (!__is.fail())
255         swap(__x.__p_.__p_, __p);
256     return __is;
257 }
258 
259 _LIBCPP_END_NAMESPACE_STD
260 
261 _LIBCPP_POP_MACROS
262 
263 #endif // _LIBCPP___RANDOM_DISCRETE_DISTRIBUTION_H
264