rand: Zipf distributed random variates.

R=rsc
https://golang.org/cl/176070

diff --git a/src/pkg/rand/Makefile b/src/pkg/rand/Makefile
index 068b6fdee0b579e993d37baeeb060a0e41aeb79c..bf3bd87d6769df2db372ec56f673022a3ec1bdaf 100644 (file)
--- a/src/pkg/rand/Makefile
+++ b/src/pkg/rand/Makefile
@@ -10,5 +10,6 @@ GOFILES=\
        normal.go\
        rand.go\
        rng.go\
+       zipf.go\
 
 include ../../Make.pkg

diff --git a/src/pkg/rand/zipf.go b/src/pkg/rand/zipf.go
new file mode 100644 (file)
index 0000000..c4e7b7d
--- /dev/null
+++ b/src/pkg/rand/zipf.go
@@ -0,0 +1,73 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// W.Hormann, G.Derflinger:
+// "Rejection-Inversion to Generate Variates
+// from Monotone Discrete Distributions"
+// http://eeyore.wu-wien.ac.at/papers/96-04-04.wh-der.ps.gz
+
+package rand
+
+import "math"
+
+// A Zipf generates Zipf distributed variates.
+type Zipf struct {
+       r            *Rand
+       imax         float64
+       v            float64
+       q            float64
+       s            float64
+       oneminusQ    float64
+       oneminusQinv float64
+       hxm          float64
+       hx0minusHxm  float64
+}
+
+func (z *Zipf) h(x float64) float64 {
+       return math.Exp(z.oneminusQ*math.Log(z.v+x)) * z.oneminusQinv
+}
+
+func (z *Zipf) hinv(x float64) float64 {
+       return math.Exp(z.oneminusQinv*math.Log(z.oneminusQ*x)) - z.v
+}
+
+// NewZipf returns a Zipf generating variates p(k) on [0, imax]
+// proportional to (v+k)**(-s) where s>1 and k>=0, and v>=1.
+//
+func NewZipf(r *Rand, s float64, v float64, imax uint64) *Zipf {
+       z := new(Zipf)
+       if s <= 1.0 || v < 1 {
+               return nil
+       }
+       z.r = r
+       z.imax = float64(imax)
+       z.v = v
+       z.q = s
+       z.oneminusQ = 1.0 - z.q
+       z.oneminusQinv = 1.0 / z.oneminusQ
+       z.hxm = z.h(z.imax + 0.5)
+       z.hx0minusHxm = z.h(0.5) - math.Exp(math.Log(z.v)*(-z.q)) - z.hxm
+       z.s = 1 - z.hinv(z.h(1.5)-math.Exp(-z.q*math.Log(z.v+1.0)))
+       return z
+}
+
+// Uint64 returns a value drawn from the Zipf distributed described
+// by the Zipf object.
+func (z *Zipf) Uint64() uint64 {
+       k := float64(0.0)
+
+       for {
+               r := z.r.Float64() // r on [0,1]
+               ur := z.hxm + r*z.hx0minusHxm
+               x := z.hinv(ur)
+               k = math.Floor(x + 0.5)
+               if k-x <= z.s {
+                       break
+               }
+               if ur >= z.h(k+0.5)-math.Exp(-math.Log(k+z.v)*z.q) {
+                       break
+               }
+       }
+       return uint64(k)
+}