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 (
       8             :         "math/rand/v2"
       9             :         "sync"
      10             : )
      11             : 
      12             : // Deck is a random number generator that generates numbers in the range
      13             : // [0,len(weights)-1] where the probability of i is
      14             : // weights(i)/sum(weights). Unlike Weighted, the weights are specified as
      15             : // integers and used in a deck-of-cards style random number selection which
      16             : // ensures that each element is returned with a desired frequency within the
      17             : // size of the deck.
      18             : type Deck struct {
      19             :         rng *rand.Rand
      20             :         mu  struct {
      21             :                 sync.Mutex
      22             :                 index int
      23             :                 deck  []int
      24             :         }
      25             : }
      26             : 
      27             : // NewDeck returns a new deck random number generator.
      28           0 : func NewDeck(rng *rand.Rand, weights ...int) *Deck {
      29           0 :         var sum int
      30           0 :         for i := range weights {
      31           0 :                 sum += weights[i]
      32           0 :         }
      33           0 :         deck := make([]int, 0, sum)
      34           0 :         for i := range weights {
      35           0 :                 for j := 0; j < weights[i]; j++ {
      36           0 :                         deck = append(deck, i)
      37           0 :                 }
      38             :         }
      39           0 :         d := &Deck{
      40           0 :                 rng: ensureRand(rng),
      41           0 :         }
      42           0 :         d.mu.index = len(deck)
      43           0 :         d.mu.deck = deck
      44           0 :         return d
      45             : }
      46             : 
      47             : // Int returns a random number in the range [0,len(weights)-1] where the
      48             : // probability of i is weights(i)/sum(weights).
      49           0 : func (d *Deck) Int() int {
      50           0 :         d.mu.Lock()
      51           0 :         if d.mu.index == len(d.mu.deck) {
      52           0 :                 d.rng.Shuffle(len(d.mu.deck), func(i, j int) {
      53           0 :                         d.mu.deck[i], d.mu.deck[j] = d.mu.deck[j], d.mu.deck[i]
      54           0 :                 })
      55           0 :                 d.mu.index = 0
      56             :         }
      57           0 :         result := d.mu.deck[d.mu.index]
      58           0 :         d.mu.index++
      59           0 :         d.mu.Unlock()
      60           0 :         return result
      61             : }