zipf_distribution.h (9213B)
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 15 #ifndef ABSL_RANDOM_ZIPF_DISTRIBUTION_H_ 16 #define ABSL_RANDOM_ZIPF_DISTRIBUTION_H_ 17 18 #include <cassert> 19 #include <cmath> 20 #include <istream> 21 #include <limits> 22 #include <ostream> 23 #include <type_traits> 24 25 #include "absl/base/config.h" 26 #include "absl/random/internal/iostream_state_saver.h" 27 #include "absl/random/internal/traits.h" 28 #include "absl/random/uniform_real_distribution.h" 29 30 namespace absl { 31 ABSL_NAMESPACE_BEGIN 32 33 // absl::zipf_distribution produces random integer-values in the range [0, k], 34 // distributed according to the unnormalized discrete probability function: 35 // 36 // P(x) = (v + x) ^ -q 37 // 38 // The parameter `v` must be greater than 0 and the parameter `q` must be 39 // greater than 1. If either of these parameters take invalid values then the 40 // behavior is undefined. 41 // 42 // IntType is the result_type generated by the generator. It must be of integral 43 // type; a static_assert ensures this is the case. 44 // 45 // The implementation is based on W.Hormann, G.Derflinger: 46 // 47 // "Rejection-Inversion to Generate Variates from Monotone Discrete 48 // Distributions" 49 // 50 // http://eeyore.wu-wien.ac.at/papers/96-04-04.wh-der.ps.gz 51 // 52 template <typename IntType = int> 53 class zipf_distribution { 54 public: 55 using result_type = IntType; 56 57 class param_type { 58 public: 59 using distribution_type = zipf_distribution; 60 61 // Preconditions: k >= 0, v > 0, q > 1 62 // The preconditions are validated when NDEBUG is not defined via 63 // a pair of assert() directives. 64 // If NDEBUG is defined and either or both of these parameters take invalid 65 // values, the behavior of the class is undefined. 66 explicit param_type(result_type k = (std::numeric_limits<IntType>::max)(), 67 double q = 2.0, double v = 1.0); 68 69 result_type k() const { return k_; } 70 double q() const { return q_; } 71 double v() const { return v_; } 72 73 friend bool operator==(const param_type& a, const param_type& b) { 74 return a.k_ == b.k_ && a.q_ == b.q_ && a.v_ == b.v_; 75 } 76 friend bool operator!=(const param_type& a, const param_type& b) { 77 return !(a == b); 78 } 79 80 private: 81 friend class zipf_distribution; 82 inline double h(double x) const; 83 inline double hinv(double x) const; 84 inline double compute_s() const; 85 inline double pow_negative_q(double x) const; 86 87 // Parameters here are exactly the same as the parameters of Algorithm ZRI 88 // in the paper. 89 IntType k_; 90 double q_; 91 double v_; 92 93 double one_minus_q_; // 1-q 94 double s_; 95 double one_minus_q_inv_; // 1 / 1-q 96 double hxm_; // h(k + 0.5) 97 double hx0_minus_hxm_; // h(x0) - h(k + 0.5) 98 99 static_assert(random_internal::IsIntegral<IntType>::value, 100 "Class-template absl::zipf_distribution<> must be " 101 "parameterized using an integral type."); 102 }; 103 104 zipf_distribution() 105 : zipf_distribution((std::numeric_limits<IntType>::max)()) {} 106 107 explicit zipf_distribution(result_type k, double q = 2.0, double v = 1.0) 108 : param_(k, q, v) {} 109 110 explicit zipf_distribution(const param_type& p) : param_(p) {} 111 112 void reset() {} 113 114 template <typename URBG> 115 result_type operator()(URBG& g) { // NOLINT(runtime/references) 116 return (*this)(g, param_); 117 } 118 119 template <typename URBG> 120 result_type operator()(URBG& g, // NOLINT(runtime/references) 121 const param_type& p); 122 123 result_type k() const { return param_.k(); } 124 double q() const { return param_.q(); } 125 double v() const { return param_.v(); } 126 127 param_type param() const { return param_; } 128 void param(const param_type& p) { param_ = p; } 129 130 result_type(min)() const { return 0; } 131 result_type(max)() const { return k(); } 132 133 friend bool operator==(const zipf_distribution& a, 134 const zipf_distribution& b) { 135 return a.param_ == b.param_; 136 } 137 friend bool operator!=(const zipf_distribution& a, 138 const zipf_distribution& b) { 139 return a.param_ != b.param_; 140 } 141 142 private: 143 param_type param_; 144 }; 145 146 // -------------------------------------------------------------------------- 147 // Implementation details follow 148 // -------------------------------------------------------------------------- 149 150 template <typename IntType> 151 zipf_distribution<IntType>::param_type::param_type( 152 typename zipf_distribution<IntType>::result_type k, double q, double v) 153 : k_(k), q_(q), v_(v), one_minus_q_(1 - q) { 154 assert(q > 1); 155 assert(v > 0); 156 assert(k >= 0); 157 one_minus_q_inv_ = 1 / one_minus_q_; 158 159 // Setup for the ZRI algorithm (pg 17 of the paper). 160 // Compute: h(i max) => h(k + 0.5) 161 constexpr double kMax = 18446744073709549568.0; 162 double kd = static_cast<double>(k); 163 // TODO(absl-team): Determine if this check is needed, and if so, add a test 164 // that fails for k > kMax 165 if (kd > kMax) { 166 // Ensure that our maximum value is capped to a value which will 167 // round-trip back through double. 168 kd = kMax; 169 } 170 hxm_ = h(kd + 0.5); 171 172 // Compute: h(0) 173 const bool use_precomputed = (v == 1.0 && q == 2.0); 174 const double h0x5 = use_precomputed ? (-1.0 / 1.5) // exp(-log(1.5)) 175 : h(0.5); 176 const double elogv_q = (v_ == 1.0) ? 1 : pow_negative_q(v_); 177 178 // h(0) = h(0.5) - exp(log(v) * -q) 179 hx0_minus_hxm_ = (h0x5 - elogv_q) - hxm_; 180 181 // And s 182 s_ = use_precomputed ? 0.46153846153846123 : compute_s(); 183 } 184 185 template <typename IntType> 186 double zipf_distribution<IntType>::param_type::h(double x) const { 187 // std::exp(one_minus_q_ * std::log(v_ + x)) * one_minus_q_inv_; 188 x += v_; 189 return (one_minus_q_ == -1.0) 190 ? (-1.0 / x) // -exp(-log(x)) 191 : (std::exp(std::log(x) * one_minus_q_) * one_minus_q_inv_); 192 } 193 194 template <typename IntType> 195 double zipf_distribution<IntType>::param_type::hinv(double x) const { 196 // std::exp(one_minus_q_inv_ * std::log(one_minus_q_ * x)) - v_; 197 return -v_ + ((one_minus_q_ == -1.0) 198 ? (-1.0 / x) // exp(-log(-x)) 199 : std::exp(one_minus_q_inv_ * std::log(one_minus_q_ * x))); 200 } 201 202 template <typename IntType> 203 double zipf_distribution<IntType>::param_type::compute_s() const { 204 // 1 - hinv(h(1.5) - std::exp(std::log(v_ + 1) * -q_)); 205 return 1.0 - hinv(h(1.5) - pow_negative_q(v_ + 1.0)); 206 } 207 208 template <typename IntType> 209 double zipf_distribution<IntType>::param_type::pow_negative_q(double x) const { 210 // std::exp(std::log(x) * -q_); 211 return q_ == 2.0 ? (1.0 / (x * x)) : std::exp(std::log(x) * -q_); 212 } 213 214 template <typename IntType> 215 template <typename URBG> 216 typename zipf_distribution<IntType>::result_type 217 zipf_distribution<IntType>::operator()( 218 URBG& g, const param_type& p) { // NOLINT(runtime/references) 219 absl::uniform_real_distribution<double> uniform_double; 220 double k; 221 for (;;) { 222 const double v = uniform_double(g); 223 const double u = p.hxm_ + v * p.hx0_minus_hxm_; 224 const double x = p.hinv(u); 225 k = rint(x); // std::floor(x + 0.5); 226 if (k > static_cast<double>(p.k())) continue; // reject k > max_k 227 if (k - x <= p.s_) break; 228 const double h = p.h(k + 0.5); 229 const double r = p.pow_negative_q(p.v_ + k); 230 if (u >= h - r) break; 231 } 232 IntType ki = static_cast<IntType>(k); 233 assert(ki <= p.k_); 234 return ki; 235 } 236 237 template <typename CharT, typename Traits, typename IntType> 238 std::basic_ostream<CharT, Traits>& operator<<( 239 std::basic_ostream<CharT, Traits>& os, // NOLINT(runtime/references) 240 const zipf_distribution<IntType>& x) { 241 using stream_type = 242 typename random_internal::stream_format_type<IntType>::type; 243 auto saver = random_internal::make_ostream_state_saver(os); 244 os.precision(random_internal::stream_precision_helper<double>::kPrecision); 245 os << static_cast<stream_type>(x.k()) << os.fill() << x.q() << os.fill() 246 << x.v(); 247 return os; 248 } 249 250 template <typename CharT, typename Traits, typename IntType> 251 std::basic_istream<CharT, Traits>& operator>>( 252 std::basic_istream<CharT, Traits>& is, // NOLINT(runtime/references) 253 zipf_distribution<IntType>& x) { // NOLINT(runtime/references) 254 using result_type = typename zipf_distribution<IntType>::result_type; 255 using param_type = typename zipf_distribution<IntType>::param_type; 256 using stream_type = 257 typename random_internal::stream_format_type<IntType>::type; 258 stream_type k; 259 double q; 260 double v; 261 262 auto saver = random_internal::make_istream_state_saver(is); 263 is >> k >> q >> v; 264 if (!is.fail()) { 265 x.param(param_type(static_cast<result_type>(k), q, v)); 266 } 267 return is; 268 } 269 270 ABSL_NAMESPACE_END 271 } // namespace absl 272 273 #endif // ABSL_RANDOM_ZIPF_DISTRIBUTION_H_