53 lines
883 B
Go
53 lines
883 B
Go
package selector
|
|
|
|
import (
|
|
"math/rand"
|
|
"time"
|
|
)
|
|
|
|
type randomWeightedItem[T any] struct {
|
|
item T
|
|
weight int
|
|
}
|
|
|
|
type randomWeighted[T any] struct {
|
|
items []*randomWeightedItem[T]
|
|
sum int
|
|
r *rand.Rand
|
|
}
|
|
|
|
func newRandomWeighted[T any]() *randomWeighted[T] {
|
|
return &randomWeighted[T]{
|
|
r: rand.New(rand.NewSource(time.Now().UnixNano())),
|
|
}
|
|
}
|
|
|
|
func (rw *randomWeighted[T]) Add(item T, weight int) {
|
|
ri := &randomWeightedItem[T]{item: item, weight: weight}
|
|
rw.items = append(rw.items, ri)
|
|
rw.sum += weight
|
|
}
|
|
|
|
func (rw *randomWeighted[T]) Next() (v T) {
|
|
if len(rw.items) == 0 {
|
|
return
|
|
}
|
|
if rw.sum <= 0 {
|
|
return
|
|
}
|
|
weight := rw.r.Intn(rw.sum) + 1
|
|
for _, item := range rw.items {
|
|
weight -= item.weight
|
|
if weight <= 0 {
|
|
return item.item
|
|
}
|
|
}
|
|
|
|
return rw.items[len(rw.items)-1].item
|
|
}
|
|
|
|
func (rw *randomWeighted[T]) Reset() {
|
|
rw.items = nil
|
|
rw.sum = 0
|
|
}
|