init

go-gost/x/selector/weighted.go

@@ -0,0 +1,52 @@
+package selector
+
+import (
+	"math/rand"
+	"time"
+)
+
+type randomWeightedItem[T any] struct {
+	item   T
+	weight int
+}
+
+type RandomWeighted[T any] struct {
+	items []*randomWeightedItem[T]
+	sum   int
+	r     *rand.Rand
+}
+
+func NewRandomWeighted[T any]() *RandomWeighted[T] {
+	return &RandomWeighted[T]{
+		r: rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+}
+
+func (rw *RandomWeighted[T]) Add(item T, weight int) {
+	ri := &randomWeightedItem[T]{item: item, weight: weight}
+	rw.items = append(rw.items, ri)
+	rw.sum += weight
+}
+
+func (rw *RandomWeighted[T]) Next() (v T) {
+	if len(rw.items) == 0 {
+		return
+	}
+	if rw.sum <= 0 {
+		return
+	}
+	weight := rw.r.Intn(rw.sum) + 1
+	for _, item := range rw.items {
+		weight -= item.weight
+		if weight <= 0 {
+			return item.item
+		}
+	}
+
+	return rw.items[len(rw.items)-1].item
+}
+
+func (rw *RandomWeighted[T]) Reset() {
+	rw.items = nil
+	rw.sum = 0
+}