5 changed files with 80 additions and 1 deletions
@ -0,0 +1,37 @@
|
||||
package rands |
||||
|
||||
import ( |
||||
"math/rand/v2" |
||||
"time" |
||||
) |
||||
|
||||
var ( |
||||
rad = rand.New(rand.NewPCG(uint64(time.Now().Unix()), uint64(time.Now().UnixNano()))) |
||||
) |
||||
|
||||
// Choice 一个选项包含一个泛型项以及一个权重,该权重用于控制其被选中的频率
|
||||
type Choice[W int | int32 | int64, T any] struct { |
||||
Weight W |
||||
Item T |
||||
} |
||||
|
||||
// WeightedChoice 加权随机选择, 从所提供的选项中返回一个选项, 权重值为 0 的选项永远不会被选中
|
||||
// Based on this algorithm: http://eli.thegreenplace.net/2010/01/22/weighted-random-generation-in-python/
|
||||
func WeightedChoice[W int | int32 | int64, T any](choices []Choice[W, T], rd ...*rand.Rand) (ret Choice[W, T], ok bool) { |
||||
sum := int64(0) |
||||
for _, c := range choices { |
||||
sum += int64(c.Weight) |
||||
} |
||||
crd := rad |
||||
if len(rd) > 0 && rd[0] != nil { |
||||
crd = rd[0] |
||||
} |
||||
r := crd.Int64N(sum) |
||||
for _, c := range choices { |
||||
r -= int64(c.Weight) |
||||
if r < 0 { |
||||
return c, true |
||||
} |
||||
} |
||||
return |
||||
} |
||||
@ -0,0 +1,24 @@
|
||||
package rands |
||||
|
||||
import ( |
||||
"fmt" |
||||
"testing" |
||||
) |
||||
|
||||
func TestWeightedChoice(t *testing.T) { |
||||
choices := []Choice[int, int]{ |
||||
{Weight: 10, Item: 300}, |
||||
{Weight: 1, Item: 100}, |
||||
{Weight: 5, Item: 200}, |
||||
} |
||||
times := make(map[int]int) |
||||
for range 100000 { |
||||
c, ok := WeightedChoice(choices) |
||||
if !ok { |
||||
t.Error("weight choice error") |
||||
return |
||||
} |
||||
times[c.Item]++ |
||||
} |
||||
fmt.Println(times) // map[100:6156 200:31280 300:62564]
|
||||
} |
||||
Loading…
Reference in new issue