github.com/tcnksm/go@v0.0.0-20141208075154-439b32936367/src/math/rand/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 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 generating variates p(k) on [0, imax]
    36  // proportional to (v+k)**(-s) where s>1 and k>=0, and v>=1.
    37  func NewZipf(r *Rand, s float64, v float64, imax uint64) *Zipf {
    38  	z := new(Zipf)
    39  	if s <= 1.0 || v < 1 {
    40  		return nil
    41  	}
    42  	z.r = r
    43  	z.imax = float64(imax)
    44  	z.v = v
    45  	z.q = s
    46  	z.oneminusQ = 1.0 - z.q
    47  	z.oneminusQinv = 1.0 / z.oneminusQ
    48  	z.hxm = z.h(z.imax + 0.5)
    49  	z.hx0minusHxm = z.h(0.5) - math.Exp(math.Log(z.v)*(-z.q)) - z.hxm
    50  	z.s = 1 - z.hinv(z.h(1.5)-math.Exp(-z.q*math.Log(z.v+1.0)))
    51  	return z
    52  }
    53  
    54  // Uint64 returns a value drawn from the Zipf distribution described
    55  // by the Zipf object.
    56  func (z *Zipf) Uint64() uint64 {
    57  	if z == nil {
    58  		panic("rand: nil Zipf")
    59  	}
    60  	k := 0.0
    61  
    62  	for {
    63  		r := z.r.Float64() // r on [0,1]
    64  		ur := z.hxm + r*z.hx0minusHxm
    65  		x := z.hinv(ur)
    66  		k = math.Floor(x + 0.5)
    67  		if k-x <= z.s {
    68  			break
    69  		}
    70  		if ur >= z.h(k+0.5)-math.Exp(-math.Log(k+z.v)*z.q) {
    71  			break
    72  		}
    73  	}
    74  	return uint64(k)
    75  }