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