s_cherry_e/sr/internal/extend/rand/weight.go

package rand

import "math/rand"

type (
	Weight[V any] struct {
		items      []*item[V]
		totalScore int32
	}

	item[V any] struct {
		value V
		score int32
	}
)

func NewSliceWeight[V any](list []V, fn func(V) int32) *Weight[V] {
	weight := &Weight[V]{}

	for _, v := range list {
		score := fn(v)
		weight.addItem(v, score)
	}

	return weight
}

func NewMapWeight[K comparable, V any](maps map[K]V, fn func(V) int32) *Weight[V] {
	weight := &Weight[V]{}

	for _, v := range maps {
		score := fn(v)
		weight.addItem(v, score)
	}

	return weight
}

func (p *Weight[V]) addItem(v V, score int32) {
	if score < 1 {
		// pass
		return
	}

	p.totalScore += score
	p.items = append(p.items, &item[V]{
		value: v,
		score: p.totalScore,
	})
}

func (p *Weight[V]) SelectOne() (v V) {
	if p.totalScore < 1 {
		return
	}

	randScore := rand.Int31n(p.totalScore)
	for _, row := range p.items {
		if row.score >= randScore {
			return row.value
		}
	}

	return
}

func (p *Weight[V]) Select(count int) []V {
	if p.totalScore < 1 {
		return nil
	}

	s := make([]V, count)
	for i := range count {
		randScore := rand.Int31n(p.totalScore)
		for _, row := range p.items {
			if row.score >= randScore {
				s[i] = row.value
				break
			}
		}
	}

	return s
}

func (p *Weight[V]) All() []V {
	s := make([]V, len(p.items))
	for i, row := range p.items {
		s[i] = row.value
	}
	return s
}

func (p *Weight[V]) SelectIfPresent() (v V, ok bool) {
	if p.totalScore < 1 {
		return
	}

	randScore := rand.Int31n(p.totalScore)
	for _, row := range p.items {
		if row.score >= randScore {
			return row.value, true
		}
	}

	return
}

func (p *Weight[V]) TotalScore() int32 {
	return p.totalScore
}

func SelectMapping[V any, I any](p *Weight[V], num int32, mapper func(V) (I, bool)) []I {
	if mapper == nil {
		return []I{}
	}

	s := []I{}
	for i := int32(0); i < num; i++ {
		if mapValue, ok := mapper(p.SelectOne()); ok {
			s = append(s, mapValue)
		}
	}
	return s
}

// SelectUnique 选择不重复的元素，当请求数量大于可用元素数量时返回所有元素
func (p *Weight[V]) SelectUnique(count int) []V {
	if p.totalScore < 1 {
		return nil
	}

	item_count := len(p.items)
	if count <= 0 {
		return []V{}
	}

	// 如果请求数量大于等于元素数量，直接返回所有元素
	if count >= item_count {
		return p.All()
	}

	// 创建副本以避免修改原始数据
	remainingItems := make([]*item[V], item_count)
	copy(remainingItems, p.items)

	result := make([]V, 0, count)

	for i := 0; i < count && len(remainingItems) > 0; i++ {
		remainingTotal := calculateTotalScore(remainingItems)
		randScore := rand.Int31n(remainingTotal)

		// 在剩余元素中查找
		selectedIndex := -1
		for idx, row := range remainingItems {
			if row.score >= randScore {
				selectedIndex = idx
				result = append(result, row.value)
				break
			}
		}

		if selectedIndex >= 0 {
			// 移除已选中的元素
			remainingItems = append(remainingItems[:selectedIndex], remainingItems[selectedIndex+1:]...)
			// 更新后续元素的分数
			updateScores(remainingItems, selectedIndex)
		}
	}

	return result
}

// 计算剩余元素的总分数
func calculateTotalScore[V any](items []*item[V]) int32 {
	if len(items) == 0 {
		return 0
	}
	return items[len(items)-1].score
}

// 更新从指定位置开始的元素分数
func updateScores[V any](items []*item[V], startIdx int) {
	var prevScore int32
	if startIdx > 0 {
		prevScore = items[startIdx-1].score
	} else {
		prevScore = 0
	}

	for i := startIdx; i < len(items); i++ {
		// 计算该项的实际分数值（与前一元素的差值）
		actualScore := items[i].score
		if i > 0 {
			actualScore -= items[i-1].score
		}

		// 更新累计分数
		prevScore += actualScore
		items[i].score = prevScore
	}
}