zipf.go

Documentation: math/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
View as plain text