58 lines
953 B
Go
58 lines
953 B
Go
package zipf
|
|
|
|
import (
|
|
"math"
|
|
"math/rand"
|
|
)
|
|
|
|
func NewZipf(r *rand.Rand, s float64, imax uint64) *Zipf {
|
|
z := &Zipf{}
|
|
z.rng = r
|
|
|
|
// convert to int for convenience, not sure any side effect
|
|
var n int = int(imax)
|
|
|
|
tmp := make([]float64, n)
|
|
for i := 1; i < n+1; i++ {
|
|
tmp[i-1] = 1.0 / math.Pow(float64(i), s)
|
|
}
|
|
|
|
zeta := make([]float64, n+1)
|
|
distMap := make([]float64, n+1)
|
|
z.prob = make([]float64, n)
|
|
|
|
zeta[0] = 0
|
|
for i := 1; i < n+1; i++ {
|
|
zeta[i] += zeta[i-1] + tmp[i-1]
|
|
}
|
|
|
|
for i := 0; i < n+1; i++ {
|
|
distMap[i] = zeta[i] / zeta[n]
|
|
}
|
|
|
|
for i := 1; i < n+1; i++ {
|
|
z.prob[i-1] = distMap[i] - distMap[i-1]
|
|
}
|
|
|
|
return z
|
|
}
|
|
|
|
type Zipf struct {
|
|
s float64
|
|
prob []float64
|
|
rng *rand.Rand
|
|
}
|
|
|
|
func (z *Zipf) Uint64() uint64 {
|
|
u := z.rng.Float64()
|
|
for i := 0; i < len(z.prob); i++ {
|
|
if u < z.prob[i] {
|
|
return uint64(i)
|
|
}
|
|
u -= z.prob[i]
|
|
}
|
|
// should not come here
|
|
// return last index for safty
|
|
return uint64(len(z.prob)-1)
|
|
}
|