github.com/sdqri/sequined@v0.0.0-20240421190656-fc6bf956f4d8/internal/graphgenerator/utils.go (about)

     1  package graphgenerator
     2  
     3  import (
     4  	"errors"
     5  	"math/rand"
     6  	"time"
     7  )
     8  
     9  var (
    10  	ErrNoProbabilities      = errors.New("no probabilities provided")
    11  	ErrZeroSumProbabilities = errors.New("sum of probabilities is zero")
    12  	ErrFailedToSelectIndex  = errors.New("failed to select index")
    13  )
    14  
    15  func SelectByProbability(probabilities []float64) (int, error) {
    16  	if len(probabilities) == 0 {
    17  		return -1, ErrNoProbabilities
    18  	}
    19  
    20  	source := rand.NewSource(time.Now().UnixNano())
    21  	rng := rand.New(source)
    22  
    23  	// Calculate the sum of probabilities
    24  	var sum float64 = 0
    25  	for _, prob := range probabilities {
    26  		sum += prob
    27  	}
    28  
    29  	if sum == 0 {
    30  		return -1, ErrZeroSumProbabilities
    31  	}
    32  
    33  	// Generate a random number between 0 and the sum of probabilities
    34  	r := rng.Float64() * sum
    35  
    36  	var cumulativeProb float64 = 0
    37  	for i, prob := range probabilities {
    38  		cumulativeProb += prob
    39  		if r < cumulativeProb {
    40  			return i, nil
    41  		}
    42  	}
    43  
    44  	return -1, ErrFailedToSelectIndex
    45  }