pgregory.net/rand@v1.0.3-0.20230808192358-a0b8ce02f4da/std_zipf.go (about) 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-go 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 10 package rand 11 12 import "math" 13 14 // A Zipf generates Zipf distributed variates. 15 type 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 27 func (z *Zipf) h(x float64) float64 { 28 return math.Exp(z.oneminusQ*math.Log(z.v+x)) * z.oneminusQinv 29 } 30 31 func (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. 39 func 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. 58 func (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 }