zipf.mx raw

   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 := &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