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
k()69 result_type k() const { return k_; }
q()70 double q() const { return q_; }
v()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
zipf_distribution()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)
param_(k,q,v)108 : param_(k, q, v) {}
109
zipf_distribution(const param_type & p)110 explicit zipf_distribution(const param_type& p) : param_(p) {}
111
reset()112 void reset() {}
113
114 template <typename URBG>
operator()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
k()123 result_type k() const { return param_.k(); }
q()124 double q() const { return param_.q(); }
v()125 double v() const { return param_.v(); }
126
param()127 param_type param() const { return param_; }
param(const param_type & p)128 void param(const param_type& p) { param_ = p; }
129
result_type(min)130 result_type(min)() const { return 0; }
result_type(max)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>
param_type(typename zipf_distribution<IntType>::result_type k,double q,double v)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>
h(double x)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>
hinv(double x)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>
compute_s()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>
pow_negative_q(double x)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
operator()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_
274