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