Line data    Source code

       1             : // Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
       2             : // of this source code is governed by a BSD-style license that can be found in
       3             : // the LICENSE file.
       4             : 
       5             : package randvar
       6             : 
       7             : import "golang.org/x/exp/rand"
       8             : 
       9             : // Weighted is a random number generator that generates numbers in the range
      10             : // [0,len(weights)-1] where the probability of i is weights(i)/sum(weights).
      11             : type Weighted struct {
      12             :         rng     *rand.Rand
      13             :         sum     float64
      14             :         weights []float64
      15             : }
      16             : 
      17             : // NewWeighted returns a new weighted random number generator.
      18           1 : func NewWeighted(rng *rand.Rand, weights ...float64) *Weighted {
      19           1 :         var sum float64
      20           1 :         for i := range weights {
      21           1 :                 sum += weights[i]
      22           1 :         }
      23           1 :         return &Weighted{
      24           1 :                 rng:     ensureRand(rng),
      25           1 :                 sum:     sum,
      26           1 :                 weights: weights,
      27           1 :         }
      28             : }
      29             : 
      30             : // Int returns a random number in the range [0,len(weights)-1] where the
      31             : // probability of i is weights(i)/sum(weights).
      32           1 : func (w *Weighted) Int() int {
      33           1 :         p := w.rng.Float64() * w.sum
      34           1 :         for i, weight := range w.weights {
      35           1 :                 if p <= weight {
      36           1 :                         return i
      37           1 :                 }
      38           1 :                 p -= weight
      39             :         }
      40           0 :         return len(w.weights) - 1
      41             : }