ceremonyclient/pebble/internal/randvar/deck.go
Cassandra Heart 2e2a1e4789
v1.2.0 (#31)
2024-01-03 01:31:42 -06:00

63 lines
1.5 KiB
Go

// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
package randvar
import (
"sync"
"golang.org/x/exp/rand"
)
// Deck is a random number generator that generates numbers in the range
// [0,len(weights)-1] where the probability of i is
// weights(i)/sum(weights). Unlike Weighted, the weights are specified as
// integers and used in a deck-of-cards style random number selection which
// ensures that each element is returned with a desired frequency within the
// size of the deck.
type Deck struct {
rng *rand.Rand
mu struct {
sync.Mutex
index int
deck []int
}
}
// NewDeck returns a new deck random number generator.
func NewDeck(rng *rand.Rand, weights ...int) *Deck {
var sum int
for i := range weights {
sum += weights[i]
}
deck := make([]int, 0, sum)
for i := range weights {
for j := 0; j < weights[i]; j++ {
deck = append(deck, i)
}
}
d := &Deck{
rng: ensureRand(rng),
}
d.mu.index = len(deck)
d.mu.deck = deck
return d
}
// Int returns a random number in the range [0,len(weights)-1] where the
// probability of i is weights(i)/sum(weights).
func (d *Deck) Int() int {
d.mu.Lock()
if d.mu.index == len(d.mu.deck) {
d.rng.Shuffle(len(d.mu.deck), func(i, j int) {
d.mu.deck[i], d.mu.deck[j] = d.mu.deck[j], d.mu.deck[i]
})
d.mu.index = 0
}
result := d.mu.deck[d.mu.index]
d.mu.index++
d.mu.Unlock()
return result
}