Source file src/pkg/math/rand/zipf.go

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