diff --git a/utils/random/probability.go b/utils/random/probability.go index 6740ecf..7180091 100644 --- a/utils/random/probability.go +++ b/utils/random/probability.go @@ -8,6 +8,13 @@ import ( // ProbabilitySlice 按概率随机从切片中产生一个数据并返回命中的对象及是否未命中 // - 当总概率小于 1 将会发生未命中的情况 func ProbabilitySlice[T any](getProbabilityHandle func(data T) float64, data ...T) (hit T, miss bool) { + hit, _, miss = ProbabilitySliceIndex(getProbabilityHandle, data...) + return hit, miss +} + +// ProbabilitySliceIndex 按概率随机从切片中产生一个数据并返回命中的对象及对象索引以及是否未命中 +// - 当总概率小于 1 将会发生未命中的情况 +func ProbabilitySliceIndex[T any](getProbabilityHandle func(data T) float64, data ...T) (hit T, index int, miss bool) { rd := rand.New(rand.NewSource(time.Now().UnixNano())) var total float64 @@ -53,10 +60,10 @@ func ProbabilitySlice[T any](getProbabilityHandle func(data T) float64, data ... } } if i >= len(data) { - return hit, true + return hit, -1, true } hit = data[i] - return hit, false + return hit, i, false } // Probability 输入一个概率,返回是否命中 diff --git a/utils/random/weight.go b/utils/random/weight.go index 7bf4ce1..e931c4d 100644 --- a/utils/random/weight.go +++ b/utils/random/weight.go @@ -2,6 +2,12 @@ package random // WeightSlice 按权重随机从切片中产生一个数据并返回 func WeightSlice[T any](getWeightHandle func(data T) int, data ...T) T { + item, _ := WeightSliceIndex(getWeightHandle, data...) + return item +} + +// WeightSliceIndex 按权重随机从切片中产生一个数据并返回数据和对应索引 +func WeightSliceIndex[T any](getWeightHandle func(data T) int, data ...T) (item T, index int) { var total int var overlayWeight []int for _, d := range data { @@ -18,17 +24,25 @@ func WeightSlice[T any](getWeightHandle func(data T) int, data ...T) T { count = h } } - return data[i] + return data[i], i } // WeightMap 按权重随机从map中产生一个数据并返回 func WeightMap[K comparable, T any](getWeightHandle func(data T) int, data map[K]T) T { + item, _ := WeightMapKey(getWeightHandle, data) + return item +} + +// WeightMapKey 按权重随机从map中产生一个数据并返回数据和对应 key +func WeightMapKey[K comparable, T any](getWeightHandle func(data T) int, data map[K]T) (item T, key K) { var total int var overlayWeight []int var dataSlice = make([]T, 0, len(data)) - for _, d := range data { + var dataKeySlice = make([]K, 0, len(data)) + for k, d := range data { total += getWeightHandle(d) dataSlice = append(dataSlice, d) + dataKeySlice = append(dataKeySlice, k) overlayWeight = append(overlayWeight, total) } var r = IntN(total) @@ -41,5 +55,5 @@ func WeightMap[K comparable, T any](getWeightHandle func(data T) int, data map[K count = h } } - return dataSlice[i] + return dataSlice[i], dataKeySlice[i] }