1

common/utils/help.go

@@ -2,6 +2,8 @@ package utils
 
 import (
 	"errors"
+	"fmt"
+	"math/rand/v2"
 
 	"github.com/gogf/gf/v2/util/gconv"
 	"github.com/gogf/gf/v2/util/grand"
@@ -32,53 +34,71 @@ func RemoveLast(s string) string {
 	return string(runes[:len(runes)-1])
 }
 
-// ************************** 函数1：核心函数（双数组+int型概率，泛型）**************************
-// randomByIntProbs 核心概率计算函数：接收任意类型元素数组 + int型概率数组，实现权重随机
-// T any：支持任意类型的元素切片，职责单一，只处理纯int概率的核心逻辑
-func RandomByIntProbs[T any](natureSet []T, probs []int) (T, error) {
+// RandomByWeight 根据整数权重随机选择元素（优化后的泛型版本）
+// 入参：
+//   - elements: 待随机的元素集合（非空）
+//   - weights: 对应元素的权重（非负整数，长度需与elements一致；长度不匹配/总权重为0时降级为等概率随机）
+//
+// 返回：
+//   - 随机选中的元素
+//   - 错误（仅当elements为空/权重值为负时返回）
+func RandomByWeight[Element any, Weight integer](elements []Element, weights []Weight) (Element, error) {
 	// 定义泛型零值，用于错误返回
-	var zeroT T
+	var zero Element
 
-	// 1. 合法性校验：元素数组为空 或 概率数组与元素数组长度不匹配
-	if len(natureSet) == 0 {
-		return zeroT, errors.New("natureSet is empty")
-	}
-	if len(probs) == 0 || len(natureSet) != len(probs) {
-		// 长度不匹配，降级为等概率随机（兼容原有逻辑）
-		return natureSet[grand.Intn(len(natureSet))], nil
+	// 1. 核心合法性校验：元素集合不能为空
+	if len(elements) == 0 {
+		return zero, errors.New("elements set is empty (cannot random from empty slice)")
 	}
 
-	// 2. 校验概率值非负，并计算总概率
-	totalProb := 0
-	for i, p := range probs {
-		if p < 0 {
-			return zeroT, errors.New("invalid prob value: index " + gconv.String(i) + " (must be non-negative integer)")
+	// 2. 权重数组合法性校验：长度不匹配/为空时，降级为等概率随机（兼容原逻辑）
+	elemLen := len(elements)
+	if len(weights) == 0 || len(weights) != elemLen {
+		return elements[rand.IntN(elemLen)], nil
+	}
+
+	// 3. 校验权重非负，并计算总权重（统一转为int64避免溢出）
+	var totalWeight int64
+	// 预转换权重为int64，避免重复类型转换
+	intWeights := make([]int64, elemLen)
+	for i, w := range weights {
+		intW := int64(w)
+		if intW < 0 {
+			return zero, fmt.Errorf("invalid negative weight at index %d (value: %d)", i, w)
 		}
-		totalProb += p
+		intWeights[i] = intW
+		totalWeight += intW
 	}
 
-	// 3. 总概率为0，降级为等概率随机
-	if totalProb == 0 {
-		return natureSet[grand.Intn(len(natureSet))], nil
+	// 4. 总权重为0时，降级为等概率随机
+	if totalWeight == 0 {
+		return elements[rand.IntN(elemLen)], nil
 	}
 
-	// 4. 计算前缀和（权重随机核心逻辑）
-	prefixSum := make([]int, len(probs))
-	prefixSum[0] = probs[0]
-	for i := 1; i < len(probs); i++ {
-		prefixSum[i] = prefixSum[i-1] + probs[i]
+	// 5. 计算前缀和（权重随机核心逻辑，复用预转换的int64权重）
+	prefixWeights := make([]int64, elemLen)
+	prefixWeights[0] = intWeights[0]
+	for i := 1; i < elemLen; i++ {
+		prefixWeights[i] = prefixWeights[i-1] + intWeights[i]
 	}
 
-	// 5. 生成随机数匹配前缀和，返回对应元素
-	randVal := grand.Intn(totalProb)
-	for i, sum := range prefixSum {
-		if randVal < sum {
-			return natureSet[i], nil
+	// 6. 生成随机数并匹配前缀和（用int64避免溢出）
+	randomValue := grand.Intn(int(totalWeight))
+	for i, sum := range prefixWeights {
+		if randomValue < int(sum) {
+			return elements[i], nil
 		}
 	}
 
-	// 极端情况兜底，返回第一个元素
-	return natureSet[0], nil
+	// 极端兜底：理论上不会走到这里（randomValue < totalWeight），返回第一个元素保证不panic
+	return elements[0], nil
+}
+
+// integer 自定义泛型约束：匹配所有整数类型（扩展原~int的限制）
+// 包含：int/int8/int16/int32/int64/uint/uint8/uint16/uint32/uint64/uintptr
+type integer interface {
+	~int | ~int8 | ~int16 | ~int32 | ~int64 |
+		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr
 }
 
 // ************************** 函数2：封装函数（兼容string型概率，泛型）**************************
@@ -89,7 +109,7 @@ func RandomByProbs[T any](natureSet []T, probs []string) (T, error) {
 
 	// 1. 若string概率数组为空，直接调用核心函数（核心函数会处理降级逻辑）
 	if len(probs) == 0 {
-		return RandomByIntProbs(natureSet, []int{})
+		return RandomByWeight(natureSet, []int{})
 	}
 
 	// 2. string[] 转换为 int[]（使用 gconv.Int 完成转换）
@@ -100,5 +120,5 @@ func RandomByProbs[T any](natureSet []T, probs []string) (T, error) {
 	}
 
 	// 3. 调用核心函数，复用概率计算逻辑
-	return RandomByIntProbs(natureSet, probInts)
+	return RandomByWeight(natureSet, probInts)
 }