fox/xrand/weight.go

package xrand

import "errors"

var ErrWeightRandomBadParam = errors.New("WeightRandom: bad param")

type randomItem[T any] struct {
	Weight uint
	Data   T
}

// WeightRandom 加权随机
func WeightRandom[W int | uint | int32 | uint32 | int64 | uint64, T any](m map[W]T) (result T, err error) {
	sum := uint(0)
	var items []randomItem[T]
	for weight, data := range m {
		sum += uint(weight)
		items = append(items, randomItem[T]{
			Weight: uint(weight),
			Data:   data,
		})
	}

	if len(items) == 0 || sum == 0 {
		err = ErrWeightRandomBadParam
		return
	}
	if len(items) == 1 {
		return items[0].Data, nil
	}

	r := IntN(int(sum))
	for _, item := range items {
		r -= int(item.Weight)
		if r < 0 {
			return item.Data, nil
		}
	}

	err = ErrWeightRandomBadParam
	return
}