github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/internal/randvar/weighted.go (about)

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