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 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