Home | History | Annotate | Download | only in rand
      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 
     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 }
     78