fix: random 包按权重产生结果更改为 int64

This commit is contained in:
kercylan98 2023-10-09 14:16:46 +08:00
parent 9435ba5ecb
commit 433ba08c75
1 changed files with 10 additions and 10 deletions

View File

@ -1,20 +1,20 @@
package random package random
// WeightSlice 按权重随机从切片中产生一个数据并返回 // WeightSlice 按权重随机从切片中产生一个数据并返回
func WeightSlice[T any](getWeightHandle func(data T) int, data ...T) T { func WeightSlice[T any](getWeightHandle func(data T) int64, data ...T) T {
item, _ := WeightSliceIndex(getWeightHandle, data...) item, _ := WeightSliceIndex(getWeightHandle, data...)
return item return item
} }
// WeightSliceIndex 按权重随机从切片中产生一个数据并返回数据和对应索引 // WeightSliceIndex 按权重随机从切片中产生一个数据并返回数据和对应索引
func WeightSliceIndex[T any](getWeightHandle func(data T) int, data ...T) (item T, index int) { func WeightSliceIndex[T any](getWeightHandle func(data T) int64, data ...T) (item T, index int) {
var total int var total int64
var overlayWeight []int var overlayWeight []int64
for _, d := range data { for _, d := range data {
total += getWeightHandle(d) total += getWeightHandle(d)
overlayWeight = append(overlayWeight, total) overlayWeight = append(overlayWeight, total)
} }
var r = IntN(total) var r = Int64(0, total)
var i, count = 0, len(overlayWeight) var i, count = 0, len(overlayWeight)
for i < count { for i < count {
h := int(uint(i+count) >> 1) h := int(uint(i+count) >> 1)
@ -28,15 +28,15 @@ func WeightSliceIndex[T any](getWeightHandle func(data T) int, data ...T) (item
} }
// WeightMap 按权重随机从map中产生一个数据并返回 // WeightMap 按权重随机从map中产生一个数据并返回
func WeightMap[K comparable, T any](getWeightHandle func(data T) int, data map[K]T) T { func WeightMap[K comparable, T any](getWeightHandle func(data T) int64, data map[K]T) T {
item, _ := WeightMapKey(getWeightHandle, data) item, _ := WeightMapKey(getWeightHandle, data)
return item return item
} }
// WeightMapKey 按权重随机从map中产生一个数据并返回数据和对应 key // WeightMapKey 按权重随机从map中产生一个数据并返回数据和对应 key
func WeightMapKey[K comparable, T any](getWeightHandle func(data T) int, data map[K]T) (item T, key K) { func WeightMapKey[K comparable, T any](getWeightHandle func(data T) int64, data map[K]T) (item T, key K) {
var total int var total int64
var overlayWeight []int var overlayWeight []int64
var dataSlice = make([]T, 0, len(data)) var dataSlice = make([]T, 0, len(data))
var dataKeySlice = make([]K, 0, len(data)) var dataKeySlice = make([]K, 0, len(data))
for k, d := range data { for k, d := range data {
@ -45,7 +45,7 @@ func WeightMapKey[K comparable, T any](getWeightHandle func(data T) int, data ma
dataKeySlice = append(dataKeySlice, k) dataKeySlice = append(dataKeySlice, k)
overlayWeight = append(overlayWeight, total) overlayWeight = append(overlayWeight, total)
} }
var r = IntN(total) var r = Int64(0, total)
var i, count = 0, len(overlayWeight) var i, count = 0, len(overlayWeight)
for i < count { for i < count {
h := int(uint(i+count) >> 1) h := int(uint(i+count) >> 1)