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