admin-entrance/server/package/graveserver/leafstalk/otherutils/weightrand/weightrand.go

package weightrand

import (
	"errors"
	"fmt"
	"math/rand"
	"slices"
)

// WeightedItem 结构体，用于存储项目和其权重
type WeightedItem[T comparable] struct {
	Item   T
	Weight int64
}

// WeightedRandomSelector 结构体，用于根据权重随机选择一项
type WeightedRandomSelector[T comparable] struct {
	Items []*WeightedItem[T]
}

// NewWeightedRandomSelector 创建一个新的 WeightedRandomSelector 实例
func NewWeightedRandomSelector[T comparable](cnt int) *WeightedRandomSelector[T] {
	wrs := &WeightedRandomSelector[T]{
		Items: make([]*WeightedItem[T], 0, cnt),
	}
	return wrs
}

// AddItem 添加一个新的权重项
func (wrs *WeightedRandomSelector[T]) AddItem(item T, weight int64) {
	if weight <= 0 {
		return
	}
	it := new(WeightedItem[T])
	it.Item = item
	it.Weight = weight
	wrs.Items = append(wrs.Items, it)
}

// RemoveItem 删除指定的权重项
func (wrs *WeightedRandomSelector[T]) RemoveItem(item T) {
	wrs.Items = slices.DeleteFunc(wrs.Items, func(wi *WeightedItem[T]) bool {
		return wi.Item == item
	})
}

// SelectRandom 根据权重随机选择一项
func (wrs *WeightedRandomSelector[T]) SelectRandom() (T, error) {
	var t T
	if len(wrs.Items) == 0 {
		return t, errors.New("empty items")
	}

	// 计算总权重
	totalWeight := int64(0)
	for _, item := range wrs.Items {
		totalWeight += item.Weight
	}
	if totalWeight <= 0 {
		return t, errors.New("total weight is zero")
	}

	// 生成一个 [0, totalWeight) 之间的随机数
	randomValue := rand.Int63n(totalWeight)

	// 根据随机数找到对应的项目
	for _, item := range wrs.Items {
		if randomValue < item.Weight {
			return item.Item, nil
		}
		randomValue -= item.Weight
	}

	// 理论上不会到达这里，因为总权重一定会匹配到某个项目
	return t, errors.New("unreachable code")
}

func test() {
	items := []WeightedItem[string]{
		{"Item1", 10},
		{"Item2", 20},
		{"Item3", 30},
		{"Item4", 40},
	}

	selector := NewWeightedRandomSelector[string](len(items))
	for i := 0; i < 10; i++ {
		fmt.Println(selector.SelectRandom())
	}
}