1// Copyright 2009 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// W.Hormann, G.Derflinger:
6// "Rejection-Inversion to Generate Variates
7// from Monotone Discrete Distributions"
8// http://eeyore.wu-wien.ac.at/papers/96-04-04.wh-der.ps.gz
9
10package rand
11
12import "math"
13
14// A Zipf generates Zipf distributed variates.
15type Zipf struct {
16	r            *Rand
17	imax         float64
18	v            float64
19	q            float64
20	s            float64
21	oneminusQ    float64
22	oneminusQinv float64
23	hxm          float64
24	hx0minusHxm  float64
25}
26
27func (z *Zipf) h(x float64) float64 {
28	return math.Exp(z.oneminusQ*math.Log(z.v+x)) * z.oneminusQinv
29}
30
31func (z *Zipf) hinv(x float64) float64 {
32	return math.Exp(z.oneminusQinv*math.Log(z.oneminusQ*x)) - z.v
33}
34
35// NewZipf returns a Zipf variate generator.
36// The generator generates values k ∈ [0, imax]
37// such that P(k) is proportional to (v + k) ** (-s).
38// Requirements: s > 1 and v >= 1.
39func NewZipf(r *Rand, s float64, v float64, imax uint64) *Zipf {
40	z := new(Zipf)
41	if s <= 1.0 || v < 1 {
42		return nil
43	}
44	z.r = r
45	z.imax = float64(imax)
46	z.v = v
47	z.q = s
48	z.oneminusQ = 1.0 - z.q
49	z.oneminusQinv = 1.0 / z.oneminusQ
50	z.hxm = z.h(z.imax + 0.5)
51	z.hx0minusHxm = z.h(0.5) - math.Exp(math.Log(z.v)*(-z.q)) - z.hxm
52	z.s = 1 - z.hinv(z.h(1.5)-math.Exp(-z.q*math.Log(z.v+1.0)))
53	return z
54}
55
56// Uint64 returns a value drawn from the Zipf distribution described
57// by the Zipf object.
58func (z *Zipf) Uint64() uint64 {
59	if z == nil {
60		panic("rand: nil Zipf")
61	}
62	k := 0.0
63
64	for {
65		r := z.r.Float64() // r on [0,1]
66		ur := z.hxm + r*z.hx0minusHxm
67		x := z.hinv(ur)
68		k = math.Floor(x + 0.5)
69		if k-x <= z.s {
70			break
71		}
72		if ur >= z.h(k+0.5)-math.Exp(-math.Log(k+z.v)*z.q) {
73			break
74		}
75	}
76	return uint64(k)
77}
78