absl/random/log_uniform_int_distribution.h

0001 // Copyright 2017 The Abseil Authors.
0002 //
0003 // Licensed under the Apache License, Version 2.0 (the "License");
0004 // you may not use this file except in compliance with the License.
0005 // You may obtain a copy of the License at
0006 //
0007 //      https://www.apache.org/licenses/LICENSE-2.0
0008 //
0009 // Unless required by applicable law or agreed to in writing, software
0010 // distributed under the License is distributed on an "AS IS" BASIS,
0011 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0012 // See the License for the specific language governing permissions and
0013 // limitations under the License.
0014
0015 #ifndef ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_
0016 #define ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_
0017
0018 #include <algorithm>
0019 #include <cassert>
0020 #include <cmath>
0021 #include <istream>
0022 #include <limits>
0023 #include <ostream>
0024 #include <type_traits>
0025
0026 #include "absl/numeric/bits.h"
0027 #include "absl/random/internal/fastmath.h"
0028 #include "absl/random/internal/generate_real.h"
0029 #include "absl/random/internal/iostream_state_saver.h"
0030 #include "absl/random/internal/traits.h"
0031 #include "absl/random/uniform_int_distribution.h"
0032
0033 namespace absl {
0034 ABSL_NAMESPACE_BEGIN
0035
0036 // log_uniform_int_distribution:
0037 //
0038 // Returns a random variate R in range [min, max] such that
0039 // floor(log(R-min, base)) is uniformly distributed.
0040 // We ensure uniformity by discretization using the
0041 // boundary sets [0, 1, base, base * base, ... min(base*n, max)]
0042 //
0043 template <typename IntType = int>
0044 class log_uniform_int_distribution {
0045  private:
0046   using unsigned_type =
0047       typename random_internal::make_unsigned_bits<IntType>::type;
0048
0049  public:
0050   using result_type = IntType;
0051
0052   class param_type {
0053    public:
0054     using distribution_type = log_uniform_int_distribution;
0055
0056     explicit param_type(
0057         result_type min = 0,
0058         result_type max = (std::numeric_limits<result_type>::max)(),
0059         result_type base = 2)
0060         : min_(min),
0061           max_(max),
0062           base_(base),
0063           range_(static_cast<unsigned_type>(max_) -
0064                  static_cast<unsigned_type>(min_)),
0065           log_range_(0) {
0066       assert(max_ >= min_);
0067       assert(base_ > 1);
0068
0069       if (base_ == 2) {
0070         // Determine where the first set bit is on range(), giving a log2(range)
0071         // value which can be used to construct bounds.
0072         log_range_ = (std::min)(random_internal::BitWidth(range()),
0073                                 std::numeric_limits<unsigned_type>::digits);
0074       } else {
0075         // NOTE: Computing the logN(x) introduces error from 2 sources:
0076         // 1. Conversion of int to double loses precision for values >=
0077         // 2^53, which may cause some log() computations to operate on
0078         // different values.
0079         // 2. The error introduced by the division will cause the result
0080         // to differ from the expected value.
0081         //
0082         // Thus a result which should equal K may equal K +/- epsilon,
0083         // which can eliminate some values depending on where the bounds fall.
0084         const double inv_log_base = 1.0 / std::log(static_cast<double>(base_));
0085         const double log_range = std::log(static_cast<double>(range()) + 0.5);
0086         log_range_ = static_cast<int>(std::ceil(inv_log_base * log_range));
0087       }
0088     }
0089
0090     result_type(min)() const { return min_; }
0091     result_type(max)() const { return max_; }
0092     result_type base() const { return base_; }
0093
0094     friend bool operator==(const param_type& a, const param_type& b) {
0095       return a.min_ == b.min_ && a.max_ == b.max_ && a.base_ == b.base_;
0096     }
0097
0098     friend bool operator!=(const param_type& a, const param_type& b) {
0099       return !(a == b);
0100     }
0101
0102    private:
0103     friend class log_uniform_int_distribution;
0104
0105     int log_range() const { return log_range_; }
0106     unsigned_type range() const { return range_; }
0107
0108     result_type min_;
0109     result_type max_;
0110     result_type base_;
0111     unsigned_type range_;  // max - min
0112     int log_range_;        // ceil(logN(range_))
0113
0114     static_assert(random_internal::IsIntegral<IntType>::value,
0115                   "Class-template absl::log_uniform_int_distribution<> must be "
0116                   "parameterized using an integral type.");
0117   };
0118
0119   log_uniform_int_distribution() : log_uniform_int_distribution(0) {}
0120
0121   explicit log_uniform_int_distribution(
0122       result_type min,
0123       result_type max = (std::numeric_limits<result_type>::max)(),
0124       result_type base = 2)
0125       : param_(min, max, base) {}
0126
0127   explicit log_uniform_int_distribution(const param_type& p) : param_(p) {}
0128
0129   void reset() {}
0130
0131   // generating functions
0132   template <typename URBG>
0133   result_type operator()(URBG& g) {  // NOLINT(runtime/references)
0134     return (*this)(g, param_);
0135   }
0136
0137   template <typename URBG>
0138   result_type operator()(URBG& g,  // NOLINT(runtime/references)
0139                          const param_type& p) {
0140     return static_cast<result_type>((p.min)() + Generate(g, p));
0141   }
0142
0143   result_type(min)() const { return (param_.min)(); }
0144   result_type(max)() const { return (param_.max)(); }
0145   result_type base() const { return param_.base(); }
0146
0147   param_type param() const { return param_; }
0148   void param(const param_type& p) { param_ = p; }
0149
0150   friend bool operator==(const log_uniform_int_distribution& a,
0151                          const log_uniform_int_distribution& b) {
0152     return a.param_ == b.param_;
0153   }
0154   friend bool operator!=(const log_uniform_int_distribution& a,
0155                          const log_uniform_int_distribution& b) {
0156     return a.param_ != b.param_;
0157   }
0158
0159  private:
0160   // Returns a log-uniform variate in the range [0, p.range()]. The caller
0161   // should add min() to shift the result to the correct range.
0162   template <typename URNG>
0163   unsigned_type Generate(URNG& g,  // NOLINT(runtime/references)
0164                          const param_type& p);
0165
0166   param_type param_;
0167 };
0168
0169 template <typename IntType>
0170 template <typename URBG>
0171 typename log_uniform_int_distribution<IntType>::unsigned_type
0172 log_uniform_int_distribution<IntType>::Generate(
0173     URBG& g,  // NOLINT(runtime/references)
0174     const param_type& p) {
0175   // sample e over [0, log_range]. Map the results of e to this:
0176   // 0 => 0
0177   // 1 => [1, b-1]
0178   // 2 => [b, (b^2)-1]
0179   // n => [b^(n-1)..(b^n)-1]
0180   const int e = absl::uniform_int_distribution<int>(0, p.log_range())(g);
0181   if (e == 0) {
0182     return 0;
0183   }
0184   const int d = e - 1;
0185
0186   unsigned_type base_e, top_e;
0187   if (p.base() == 2) {
0188     base_e = static_cast<unsigned_type>(1) << d;
0189
0190     top_e = (e >= std::numeric_limits<unsigned_type>::digits)
0191                 ? (std::numeric_limits<unsigned_type>::max)()
0192                 : (static_cast<unsigned_type>(1) << e) - 1;
0193   } else {
0194     const double r = std::pow(static_cast<double>(p.base()), d);
0195     const double s = (r * static_cast<double>(p.base())) - 1.0;
0196
0197     base_e =
0198         (r > static_cast<double>((std::numeric_limits<unsigned_type>::max)()))
0199             ? (std::numeric_limits<unsigned_type>::max)()
0200             : static_cast<unsigned_type>(r);
0201
0202     top_e =
0203         (s > static_cast<double>((std::numeric_limits<unsigned_type>::max)()))
0204             ? (std::numeric_limits<unsigned_type>::max)()
0205             : static_cast<unsigned_type>(s);
0206   }
0207
0208   const unsigned_type lo = (base_e >= p.range()) ? p.range() : base_e;
0209   const unsigned_type hi = (top_e >= p.range()) ? p.range() : top_e;
0210
0211   // choose uniformly over [lo, hi]
0212   return absl::uniform_int_distribution<result_type>(
0213       static_cast<result_type>(lo), static_cast<result_type>(hi))(g);
0214 }
0215
0216 template <typename CharT, typename Traits, typename IntType>
0217 std::basic_ostream<CharT, Traits>& operator<<(
0218     std::basic_ostream<CharT, Traits>& os,  // NOLINT(runtime/references)
0219     const log_uniform_int_distribution<IntType>& x) {
0220   using stream_type =
0221       typename random_internal::stream_format_type<IntType>::type;
0222   auto saver = random_internal::make_ostream_state_saver(os);
0223   os << static_cast<stream_type>((x.min)()) << os.fill()
0224      << static_cast<stream_type>((x.max)()) << os.fill()
0225      << static_cast<stream_type>(x.base());
0226   return os;
0227 }
0228
0229 template <typename CharT, typename Traits, typename IntType>
0230 std::basic_istream<CharT, Traits>& operator>>(
0231     std::basic_istream<CharT, Traits>& is,       // NOLINT(runtime/references)
0232     log_uniform_int_distribution<IntType>& x) {  // NOLINT(runtime/references)
0233   using param_type = typename log_uniform_int_distribution<IntType>::param_type;
0234   using result_type =
0235       typename log_uniform_int_distribution<IntType>::result_type;
0236   using stream_type =
0237       typename random_internal::stream_format_type<IntType>::type;
0238
0239   stream_type min;
0240   stream_type max;
0241   stream_type base;
0242
0243   auto saver = random_internal::make_istream_state_saver(is);
0244   is >> min >> max >> base;
0245   if (!is.fail()) {
0246     x.param(param_type(static_cast<result_type>(min),
0247                        static_cast<result_type>(max),
0248                        static_cast<result_type>(base)));
0249   }
0250   return is;
0251 }
0252
0253 ABSL_NAMESPACE_END
0254 }  // namespace absl
0255
0256 #endif  // ABSL_RANDOM_LOG_UNIFORM_INT_DISTRIBUTION_H_