log_uniform_int_distribution.h 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257
  1. // Copyright 2017 The Abseil Authors.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // https://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #ifndef ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_
  15. #define ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_
  16. #include <algorithm>
  17. #include <cassert>
  18. #include <cmath>
  19. #include <istream>
  20. #include <limits>
  21. #include <ostream>
  22. #include <type_traits>
  23. #include "absl/numeric/bits.h"
  24. #include "absl/random/internal/fastmath.h"
  25. #include "absl/random/internal/generate_real.h"
  26. #include "absl/random/internal/iostream_state_saver.h"
  27. #include "absl/random/internal/traits.h"
  28. #include "absl/random/uniform_int_distribution.h"
  29. namespace absl {
  30. ABSL_NAMESPACE_BEGIN
  31. // log_uniform_int_distribution:
  32. //
  33. // Returns a random variate R in range [min, max] such that
  34. // floor(log(R-min, base)) is uniformly distributed.
  35. // We ensure uniformity by discretization using the
  36. // boundary sets [0, 1, base, base * base, ... min(base*n, max)]
  37. //
  38. template <typename IntType = int>
  39. class log_uniform_int_distribution {
  40. private:
  41. using unsigned_type =
  42. typename random_internal::make_unsigned_bits<IntType>::type;
  43. public:
  44. using result_type = IntType;
  45. class param_type {
  46. public:
  47. using distribution_type = log_uniform_int_distribution;
  48. explicit param_type(
  49. result_type min = 0,
  50. result_type max = (std::numeric_limits<result_type>::max)(),
  51. result_type base = 2)
  52. : min_(min),
  53. max_(max),
  54. base_(base),
  55. range_(static_cast<unsigned_type>(max_) -
  56. static_cast<unsigned_type>(min_)),
  57. log_range_(0) {
  58. assert(max_ >= min_);
  59. assert(base_ > 1);
  60. if (base_ == 2) {
  61. // Determine where the first set bit is on range(), giving a log2(range)
  62. // value which can be used to construct bounds.
  63. log_range_ =
  64. (std::min)(bit_width(range()),
  65. static_cast<unsigned_type>(
  66. std::numeric_limits<unsigned_type>::digits));
  67. } else {
  68. // NOTE: Computing the logN(x) introduces error from 2 sources:
  69. // 1. Conversion of int to double loses precision for values >=
  70. // 2^53, which may cause some log() computations to operate on
  71. // different values.
  72. // 2. The error introduced by the division will cause the result
  73. // to differ from the expected value.
  74. //
  75. // Thus a result which should equal K may equal K +/- epsilon,
  76. // which can eliminate some values depending on where the bounds fall.
  77. const double inv_log_base = 1.0 / std::log(base_);
  78. const double log_range = std::log(static_cast<double>(range()) + 0.5);
  79. log_range_ = static_cast<int>(std::ceil(inv_log_base * log_range));
  80. }
  81. }
  82. result_type(min)() const { return min_; }
  83. result_type(max)() const { return max_; }
  84. result_type base() const { return base_; }
  85. friend bool operator==(const param_type& a, const param_type& b) {
  86. return a.min_ == b.min_ && a.max_ == b.max_ && a.base_ == b.base_;
  87. }
  88. friend bool operator!=(const param_type& a, const param_type& b) {
  89. return !(a == b);
  90. }
  91. private:
  92. friend class log_uniform_int_distribution;
  93. int log_range() const { return log_range_; }
  94. unsigned_type range() const { return range_; }
  95. result_type min_;
  96. result_type max_;
  97. result_type base_;
  98. unsigned_type range_; // max - min
  99. int log_range_; // ceil(logN(range_))
  100. static_assert(std::is_integral<IntType>::value,
  101. "Class-template absl::log_uniform_int_distribution<> must be "
  102. "parameterized using an integral type.");
  103. };
  104. log_uniform_int_distribution() : log_uniform_int_distribution(0) {}
  105. explicit log_uniform_int_distribution(
  106. result_type min,
  107. result_type max = (std::numeric_limits<result_type>::max)(),
  108. result_type base = 2)
  109. : param_(min, max, base) {}
  110. explicit log_uniform_int_distribution(const param_type& p) : param_(p) {}
  111. void reset() {}
  112. // generating functions
  113. template <typename URBG>
  114. result_type operator()(URBG& g) { // NOLINT(runtime/references)
  115. return (*this)(g, param_);
  116. }
  117. template <typename URBG>
  118. result_type operator()(URBG& g, // NOLINT(runtime/references)
  119. const param_type& p) {
  120. return (p.min)() + Generate(g, p);
  121. }
  122. result_type(min)() const { return (param_.min)(); }
  123. result_type(max)() const { return (param_.max)(); }
  124. result_type base() const { return param_.base(); }
  125. param_type param() const { return param_; }
  126. void param(const param_type& p) { param_ = p; }
  127. friend bool operator==(const log_uniform_int_distribution& a,
  128. const log_uniform_int_distribution& b) {
  129. return a.param_ == b.param_;
  130. }
  131. friend bool operator!=(const log_uniform_int_distribution& a,
  132. const log_uniform_int_distribution& b) {
  133. return a.param_ != b.param_;
  134. }
  135. private:
  136. // Returns a log-uniform variate in the range [0, p.range()]. The caller
  137. // should add min() to shift the result to the correct range.
  138. template <typename URNG>
  139. unsigned_type Generate(URNG& g, // NOLINT(runtime/references)
  140. const param_type& p);
  141. param_type param_;
  142. };
  143. template <typename IntType>
  144. template <typename URBG>
  145. typename log_uniform_int_distribution<IntType>::unsigned_type
  146. log_uniform_int_distribution<IntType>::Generate(
  147. URBG& g, // NOLINT(runtime/references)
  148. const param_type& p) {
  149. // sample e over [0, log_range]. Map the results of e to this:
  150. // 0 => 0
  151. // 1 => [1, b-1]
  152. // 2 => [b, (b^2)-1]
  153. // n => [b^(n-1)..(b^n)-1]
  154. const int e = absl::uniform_int_distribution<int>(0, p.log_range())(g);
  155. if (e == 0) {
  156. return 0;
  157. }
  158. const int d = e - 1;
  159. unsigned_type base_e, top_e;
  160. if (p.base() == 2) {
  161. base_e = static_cast<unsigned_type>(1) << d;
  162. top_e = (e >= std::numeric_limits<unsigned_type>::digits)
  163. ? (std::numeric_limits<unsigned_type>::max)()
  164. : (static_cast<unsigned_type>(1) << e) - 1;
  165. } else {
  166. const double r = std::pow(p.base(), d);
  167. const double s = (r * p.base()) - 1.0;
  168. base_e =
  169. (r > static_cast<double>((std::numeric_limits<unsigned_type>::max)()))
  170. ? (std::numeric_limits<unsigned_type>::max)()
  171. : static_cast<unsigned_type>(r);
  172. top_e =
  173. (s > static_cast<double>((std::numeric_limits<unsigned_type>::max)()))
  174. ? (std::numeric_limits<unsigned_type>::max)()
  175. : static_cast<unsigned_type>(s);
  176. }
  177. const unsigned_type lo = (base_e >= p.range()) ? p.range() : base_e;
  178. const unsigned_type hi = (top_e >= p.range()) ? p.range() : top_e;
  179. // choose uniformly over [lo, hi]
  180. return absl::uniform_int_distribution<result_type>(lo, hi)(g);
  181. }
  182. template <typename CharT, typename Traits, typename IntType>
  183. std::basic_ostream<CharT, Traits>& operator<<(
  184. std::basic_ostream<CharT, Traits>& os, // NOLINT(runtime/references)
  185. const log_uniform_int_distribution<IntType>& x) {
  186. using stream_type =
  187. typename random_internal::stream_format_type<IntType>::type;
  188. auto saver = random_internal::make_ostream_state_saver(os);
  189. os << static_cast<stream_type>((x.min)()) << os.fill()
  190. << static_cast<stream_type>((x.max)()) << os.fill()
  191. << static_cast<stream_type>(x.base());
  192. return os;
  193. }
  194. template <typename CharT, typename Traits, typename IntType>
  195. std::basic_istream<CharT, Traits>& operator>>(
  196. std::basic_istream<CharT, Traits>& is, // NOLINT(runtime/references)
  197. log_uniform_int_distribution<IntType>& x) { // NOLINT(runtime/references)
  198. using param_type = typename log_uniform_int_distribution<IntType>::param_type;
  199. using result_type =
  200. typename log_uniform_int_distribution<IntType>::result_type;
  201. using stream_type =
  202. typename random_internal::stream_format_type<IntType>::type;
  203. stream_type min;
  204. stream_type max;
  205. stream_type base;
  206. auto saver = random_internal::make_istream_state_saver(is);
  207. is >> min >> max >> base;
  208. if (!is.fail()) {
  209. x.param(param_type(static_cast<result_type>(min),
  210. static_cast<result_type>(max),
  211. static_cast<result_type>(base)));
  212. }
  213. return is;
  214. }
  215. ABSL_NAMESPACE_END
  216. } // namespace absl
  217. #endif // ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_