207 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___CXX03___RANDOM_DISCRETE_DISTRIBUTION_H10#define _LIBCPP___CXX03___RANDOM_DISCRETE_DISTRIBUTION_H11 12#include <__cxx03/__algorithm/upper_bound.h>13#include <__cxx03/__config>14#include <__cxx03/__random/is_valid.h>15#include <__cxx03/__random/uniform_real_distribution.h>16#include <__cxx03/cstddef>17#include <__cxx03/iosfwd>18#include <__cxx03/numeric>19#include <__cxx03/vector>20 21#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)22# pragma GCC system_header23#endif24 25_LIBCPP_PUSH_MACROS26#include <__cxx03/__undef_macros>27 28_LIBCPP_BEGIN_NAMESPACE_STD29 30template <class _IntType = int>31class _LIBCPP_TEMPLATE_VIS discrete_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 _LIBCPP_TEMPLATE_VIS param_type {39 vector<double> __p_;40 41 public:42 typedef discrete_distribution distribution_type;43 44 _LIBCPP_HIDE_FROM_ABI param_type() {}45 template <class _InputIterator>46 _LIBCPP_HIDE_FROM_ABI param_type(_InputIterator __f, _InputIterator __l) : __p_(__f, __l) {47 __init();48 }49 template <class _UnaryOperation>50 _LIBCPP_HIDE_FROM_ABI param_type(size_t __nw, double __xmin, double __xmax, _UnaryOperation __fw);51 52 _LIBCPP_HIDE_FROM_ABI vector<double> probabilities() const;53 54 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const param_type& __x, const param_type& __y) {55 return __x.__p_ == __y.__p_;56 }57 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const param_type& __x, const param_type& __y) { return !(__x == __y); }58 59 private:60 _LIBCPP_HIDE_FROM_ABI void __init();61 62 friend class discrete_distribution;63 64 template <class _CharT, class _Traits, class _IT>65 friend basic_ostream<_CharT, _Traits>&66 operator<<(basic_ostream<_CharT, _Traits>& __os, const discrete_distribution<_IT>& __x);67 68 template <class _CharT, class _Traits, class _IT>69 friend basic_istream<_CharT, _Traits>&70 operator>>(basic_istream<_CharT, _Traits>& __is, discrete_distribution<_IT>& __x);71 };72 73private:74 param_type __p_;75 76public:77 // constructor and reset functions78 _LIBCPP_HIDE_FROM_ABI discrete_distribution() {}79 template <class _InputIterator>80 _LIBCPP_HIDE_FROM_ABI discrete_distribution(_InputIterator __f, _InputIterator __l) : __p_(__f, __l) {}81 template <class _UnaryOperation>82 _LIBCPP_HIDE_FROM_ABI discrete_distribution(size_t __nw, double __xmin, double __xmax, _UnaryOperation __fw)83 : __p_(__nw, __xmin, __xmax, __fw) {}84 _LIBCPP_HIDE_FROM_ABI explicit discrete_distribution(const param_type& __p) : __p_(__p) {}85 _LIBCPP_HIDE_FROM_ABI void reset() {}86 87 // generating functions88 template <class _URNG>89 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g) {90 return (*this)(__g, __p_);91 }92 template <class _URNG>93 _LIBCPP_HIDE_FROM_ABI result_type operator()(_URNG& __g, const param_type& __p);94 95 // property functions96 _LIBCPP_HIDE_FROM_ABI vector<double> probabilities() const { return __p_.probabilities(); }97 98 _LIBCPP_HIDE_FROM_ABI param_type param() const { return __p_; }99 _LIBCPP_HIDE_FROM_ABI void param(const param_type& __p) { __p_ = __p; }100 101 _LIBCPP_HIDE_FROM_ABI result_type min() const { return 0; }102 _LIBCPP_HIDE_FROM_ABI result_type max() const { return __p_.__p_.size(); }103 104 friend _LIBCPP_HIDE_FROM_ABI bool operator==(const discrete_distribution& __x, const discrete_distribution& __y) {105 return __x.__p_ == __y.__p_;106 }107 friend _LIBCPP_HIDE_FROM_ABI bool operator!=(const discrete_distribution& __x, const discrete_distribution& __y) {108 return !(__x == __y);109 }110 111 template <class _CharT, class _Traits, class _IT>112 friend basic_ostream<_CharT, _Traits>&113 operator<<(basic_ostream<_CharT, _Traits>& __os, const discrete_distribution<_IT>& __x);114 115 template <class _CharT, class _Traits, class _IT>116 friend basic_istream<_CharT, _Traits>&117 operator>>(basic_istream<_CharT, _Traits>& __is, discrete_distribution<_IT>& __x);118};119 120template <class _IntType>121template <class _UnaryOperation>122discrete_distribution<_IntType>::param_type::param_type(123 size_t __nw, double __xmin, double __xmax, _UnaryOperation __fw) {124 if (__nw > 1) {125 __p_.reserve(__nw - 1);126 double __d = (__xmax - __xmin) / __nw;127 double __d2 = __d / 2;128 for (size_t __k = 0; __k < __nw; ++__k)129 __p_.push_back(__fw(__xmin + __k * __d + __d2));130 __init();131 }132}133 134template <class _IntType>135void discrete_distribution<_IntType>::param_type::__init() {136 if (!__p_.empty()) {137 if (__p_.size() > 1) {138 double __s = std::accumulate(__p_.begin(), __p_.end(), 0.0);139 for (vector<double>::iterator __i = __p_.begin(), __e = __p_.end(); __i < __e; ++__i)140 *__i /= __s;141 vector<double> __t(__p_.size() - 1);142 std::partial_sum(__p_.begin(), __p_.end() - 1, __t.begin());143 swap(__p_, __t);144 } else {145 __p_.clear();146 __p_.shrink_to_fit();147 }148 }149}150 151template <class _IntType>152vector<double> discrete_distribution<_IntType>::param_type::probabilities() const {153 size_t __n = __p_.size();154 vector<double> __p(__n + 1);155 std::adjacent_difference(__p_.begin(), __p_.end(), __p.begin());156 if (__n > 0)157 __p[__n] = 1 - __p_[__n - 1];158 else159 __p[0] = 1;160 return __p;161}162 163template <class _IntType>164template <class _URNG>165_IntType discrete_distribution<_IntType>::operator()(_URNG& __g, const param_type& __p) {166 static_assert(__libcpp_random_is_valid_urng<_URNG>::value, "");167 uniform_real_distribution<double> __gen;168 return static_cast<_IntType>(std::upper_bound(__p.__p_.begin(), __p.__p_.end(), __gen(__g)) - __p.__p_.begin());169}170 171template <class _CharT, class _Traits, class _IT>172_LIBCPP_HIDE_FROM_ABI basic_ostream<_CharT, _Traits>&173operator<<(basic_ostream<_CharT, _Traits>& __os, const discrete_distribution<_IT>& __x) {174 __save_flags<_CharT, _Traits> __lx(__os);175 typedef basic_ostream<_CharT, _Traits> _OStream;176 __os.flags(_OStream::dec | _OStream::left | _OStream::fixed | _OStream::scientific);177 _CharT __sp = __os.widen(' ');178 __os.fill(__sp);179 size_t __n = __x.__p_.__p_.size();180 __os << __n;181 for (size_t __i = 0; __i < __n; ++__i)182 __os << __sp << __x.__p_.__p_[__i];183 return __os;184}185 186template <class _CharT, class _Traits, class _IT>187_LIBCPP_HIDE_FROM_ABI basic_istream<_CharT, _Traits>&188operator>>(basic_istream<_CharT, _Traits>& __is, discrete_distribution<_IT>& __x) {189 __save_flags<_CharT, _Traits> __lx(__is);190 typedef basic_istream<_CharT, _Traits> _Istream;191 __is.flags(_Istream::dec | _Istream::skipws);192 size_t __n;193 __is >> __n;194 vector<double> __p(__n);195 for (size_t __i = 0; __i < __n; ++__i)196 __is >> __p[__i];197 if (!__is.fail())198 swap(__x.__p_.__p_, __p);199 return __is;200}201 202_LIBCPP_END_NAMESPACE_STD203 204_LIBCPP_POP_MACROS205 206#endif // _LIBCPP___CXX03___RANDOM_DISCRETE_DISTRIBUTION_H207