按权重随机

utils/random/number.go

@@ -21,3 +21,11 @@ func Float64() float64 {
 func Float32() float32 {
 	return rand.Float32()
 }
+
+// IntN 返回一个0~n的整数
+func IntN(n int) int {
+	if n <= 0 {
+		return 0
+	}
+	return rand.Intn(n)
+}

utils/random/weight.go

@@ -0,0 +1,45 @@
+package random
+
+// WeightSlice 按权重随机从切片中产生一个数据并返回
+func WeightSlice[T any](getWeightHandle func(data T) int, data ...T) T {
+	var total int
+	var overlayWeight []int
+	for _, d := range data {
+		total += getWeightHandle(d)
+		overlayWeight = append(overlayWeight, total)
+	}
+	var r = IntN(total)
+	var i, count = 0, len(overlayWeight)
+	for i < count {
+		h := int(uint(i+count) >> 1)
+		if overlayWeight[h] < r {
+			i = h + 1
+		} else {
+			count = h
+		}
+	}
+	return data[i]
+}
+
+// WeightMap 按权重随机从map中产生一个数据并返回
+func WeightMap[K comparable, T any](getWeightHandle func(data T) int, data map[K]T) T {
+	var total int
+	var overlayWeight []int
+	var dataSlice = make([]T, 0, len(data))
+	for _, d := range data {
+		total += getWeightHandle(d)
+		dataSlice = append(dataSlice, d)
+		overlayWeight = append(overlayWeight, total)
+	}
+	var r = IntN(total)
+	var i, count = 0, len(overlayWeight)
+	for i < count {
+		h := int(uint(i+count) >> 1)
+		if overlayWeight[h] < r {
+			i = h + 1
+		} else {
+			count = h
+		}
+	}
+	return dataSlice[i]
+}