43 lines
776 B
Go
43 lines
776 B
Go
|
package xrand
|
||
|
|
|
||
|
|
import "errors"
|
||
|
|
|
||
|
|
var ErrWeightRandomBadParam = errors.New("WeightRandom: bad param")
|
||
|
|
|
||
|
|
type randomItem[T any] struct {
|
||
|
|
Weight uint
|
||
|
|
Data T
|
||
|
|
}
|
||
|
|
|
||
|
|
// WeightRandom 加权随机
|
||
|
|
func WeightRandom[W int | uint | int32 | uint32 | int64 | uint64, T any](m map[W]T) (result T, err error) {
|
||
|
|
sum := uint(0)
|
||
|
|
var items []randomItem[T]
|
||
|
|
for weight, data := range m {
|
||
|
|
sum += uint(weight)
|
||
|
|
items = append(items, randomItem[T]{
|
||
|
|
Weight: uint(weight),
|
||
|
|
Data: data,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
if len(items) == 0 || sum == 0 {
|
||
|
|
err = ErrWeightRandomBadParam
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if len(items) == 1 {
|
||
|
|
return items[0].Data, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
r := Intn(int(sum))
|
||
|
|
for _, item := range items {
|
||
|
|
r -= int(item.Weight)
|
||
|
|
if r < 0 {
|
||
|
|
return item.Data, nil
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
err = ErrWeightRandomBadParam
|
||
|
|
return
|
||
|
|
}
|