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           0 : func NewWeighted(rng *rand.Rand, weights ...float64) *Weighted {
      19           0 :         var sum float64
      20           0 :         for i := range weights {
      21           0 :                 sum += weights[i]
      22           0 :         }
      23           0 :         return &Weighted{
      24           0 :                 rng:     ensureRand(rng),
      25           0 :                 sum:     sum,
      26           0 :                 weights: weights,
      27           0 :         }
      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           0 : func (w *Weighted) Int() int {
      33           0 :         p := w.rng.Float64() * w.sum
      34           0 :         for i, weight := range w.weights {
      35           0 :                 if p <= weight {
      36           0 :                         return i
      37           0 :                 }
      38           0 :                 p -= weight
      39             :         }
      40           0 :         return len(w.weights) - 1
      41             : }