zipf.go 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  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. // W.Hormann, G.Derflinger:
  5. // "Rejection-Inversion to Generate Variates
  6. // from Monotone Discrete Distributions"
  7. // http://eeyore.wu-wien.ac.at/papers/96-04-04.wh-der.ps.gz
  8. package rand
  9. import "math"
  10. // A Zipf generates Zipf distributed variates.
  11. type Zipf struct {
  12. r *Rand
  13. imax float64
  14. v float64
  15. q float64
  16. s float64
  17. oneminusQ float64
  18. oneminusQinv float64
  19. hxm float64
  20. hx0minusHxm float64
  21. }
  22. func (z *Zipf) h(x float64) float64 {
  23. return math.Exp(z.oneminusQ*math.Log(z.v+x)) * z.oneminusQinv
  24. }
  25. func (z *Zipf) hinv(x float64) float64 {
  26. return math.Exp(z.oneminusQinv*math.Log(z.oneminusQ*x)) - z.v
  27. }
  28. // NewZipf returns a Zipf generating variates p(k) on [0, imax]
  29. // proportional to (v+k)**(-s) where s>1 and k>=0, and v>=1.
  30. func NewZipf(r *Rand, s float64, v float64, imax uint64) *Zipf {
  31. z := new(Zipf)
  32. if s <= 1.0 || v < 1 {
  33. return nil
  34. }
  35. z.r = r
  36. z.imax = float64(imax)
  37. z.v = v
  38. z.q = s
  39. z.oneminusQ = 1.0 - z.q
  40. z.oneminusQinv = 1.0 / z.oneminusQ
  41. z.hxm = z.h(z.imax + 0.5)
  42. z.hx0minusHxm = z.h(0.5) - math.Exp(math.Log(z.v)*(-z.q)) - z.hxm
  43. z.s = 1 - z.hinv(z.h(1.5)-math.Exp(-z.q*math.Log(z.v+1.0)))
  44. return z
  45. }
  46. // Uint64 returns a value drawn from the Zipf distribution described
  47. // by the Zipf object.
  48. func (z *Zipf) Uint64() uint64 {
  49. if z == nil {
  50. panic("rand: nil Zipf")
  51. }
  52. k := 0.0
  53. for {
  54. r := z.r.Float64() // r on [0,1]
  55. ur := z.hxm + r*z.hx0minusHxm
  56. x := z.hinv(ur)
  57. k = math.Floor(x + 0.5)
  58. if k-x <= z.s {
  59. break
  60. }
  61. if ur >= z.h(k+0.5)-math.Exp(-math.Log(k+z.v)*z.q) {
  62. break
  63. }
  64. }
  65. return uint64(k)
  66. }