ceremonyclient/pebble/internal/randvar/weighted.go

42 lines
1.0 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
package randvar
import "golang.org/x/exp/rand"
// Weighted is a random number generator that generates numbers in the range
// [0,len(weights)-1] where the probability of i is weights(i)/sum(weights).
type Weighted struct {
rng *rand.Rand
sum float64
weights []float64
}
// NewWeighted returns a new weighted random number generator.
func NewWeighted(rng *rand.Rand, weights ...float64) *Weighted {
var sum float64
for i := range weights {
sum += weights[i]
}
return &Weighted{
rng: ensureRand(rng),
sum: sum,
weights: weights,
}
}
// Int returns a random number in the range [0,len(weights)-1] where the
// probability of i is weights(i)/sum(weights).
func (w *Weighted) Int() int {
p := w.rng.Float64() * w.sum
for i, weight := range w.weights {
if p <= weight {
return i
}
p -= weight
}
return len(w.weights) - 1
}