golang Weight Rand Helper

rand.go

package utils

import (
	"crypto/rand"
	"errors"
	"math/big"
	prand "math/rand"
	"time"
)

//权重随机器

// WeightItem interface
type WeightItem interface {
	// GetWeight 当前 Item 权重
	GetWeight() int64
}

// WeightRand 权重随机器
type WeightRand struct {
	sum        int64
	weightList []WeightItem
}

// NewWeightRand New WeightRand
func NewWeightRand(items []WeightItem) *WeightRand {
	if len(items) <= 0 {
		return nil
	}
	w := &WeightRand{
		sum:        0,
		weightList: items,
	}
	w.intWeightSum()
	return w
}

// Rand 随机出结果
func (w *WeightRand) Rand() (WeightItem, error) {
	items := w.weightList
	if len(items) <= 0 || w.sum <= 0 {
		return nil, errors.New("weight Items nil")
	}

	var err error
	var num, num2 int64
	num2, err = Rand(w.sum)
	if err != nil {
		return nil, err
	}

	for _, v := range items {
		num += v.GetWeight()
		if num > num2 {
			return v, nil
		}
	}
	return nil, errors.New("no result")
}

// initSum 初始化计算所有随机项权重总和
func (w *WeightRand) intWeightSum() {
	if w.weightList == nil {
		return
	}

	for _, v := range w.weightList {
		w.sum += v.GetWeight()
	}
}

// -----------基本随机函数----------------
var globalRand = prand.New(prand.NewSource(time.Now().UnixNano() / int64(time.Now().Minute()))) //基本伪随机

// PseudoRandRange 范围伪随机
//	seed:伪随机种子, <=0:使用默认随机种子
func PseudoRandRange(seed int64, min, max int64) (int64, error) {
	if min > max {
		return 0, errors.New("min <= max")
	}
	if min == max {
		return min, nil
	}

	r, err := PseudoRand(seed, max-min)
	if err != nil {
		return 0, err
	}
	return r + min, nil
}

// PseudoRand 伪随机
//	seed:伪随机种子, <=0:使用默认随机种子
func PseudoRand(seed int64, max int64) (int64, error) {
	if max <= 0 {
		return 0, nil
	}

	var r *prand.Rand
	if seed <= 0 {
		r = prand.New(prand.NewSource(seed))
	} else {
		r = globalRand
	}
	return r.Int63n(max), nil
}

// RandRangeFloat 范围随机
//	min/max 必须为 [0~1] 的浮点数
func RandRangeFloat(min, max float64) (float64, error) {
	if min > max {
		return 0, errors.New("min <= max")
	}
	if min == max {
		return min, nil
	}

	if max > 1 || min < 0 {
		return 0, errors.New("min/max 必须为 [0~1] 的浮点数")
	}

	b := 10000.0
	r, err := Rand(int64((max - min) * b))
	if err != nil {
		return 0, err
	}
	return float64(r)/b + min, nil
}

// RandRange 范围随机
func RandRange(min, max int64) (int64, error) {
	if min > max {
		return 0, errors.New("min <= max")
	}
	if min == max {
		return min, nil
	}

	r, err := Rand(max - min)
	if err != nil {
		return 0, err
	}
	return r + min, nil
}

// Rand 随机
func Rand(max int64) (int64, error) {
	if max <= 0 {
		return 0, nil
	}

	var bigRandNum *big.Int
	var err error
	var bigSum = big.NewInt(max)
	if bigRandNum, err = rand.Int(rand.Reader, bigSum); err != nil {
		return 0, err
	}
	return bigRandNum.Int64(), nil
}

// RandByWeight 根据某权重判断是否命中
//	weight 随机权重, >=1 直接命中
func RandByWeight(weight float64) bool {
	if weight <= 0 {
		return false
	}

	if weight >= 1.0 {
		return true
	}

	w := int64(weight * 1000.0) //归置到0~1000
	r, _ := RandRange(0, 1000)
	return r < w
}

rand_test.go

package utils

import "testing"

// tWeightItem WeightItem
type tWeightItem struct {
	Weight int64
	Value  string
}

// GetWeight 当前 Item 权重
func (w *tWeightItem) GetWeight() int64 {
	return w.Weight
}

func TestRand(t *testing.T) {
	weightItems := []WeightItem{
		&tWeightItem{Weight: 1, Value: "1"},
		&tWeightItem{Weight: 1, Value: "2"},
		&tWeightItem{Weight: 2, Value: "3"},
	}

	rand := NewWeightRand(weightItems)
	for i := 0; i < 5; i++ {
		if v, err := rand.Rand(); err != nil {
			t.Errorf("rand:%#v-err:%#v", v, err)
		} else {
			t.Logf("rand:%#v-err:%#v", v, err)
		}
	}
}