fix random weighted char because %sum can never get last value of sum

master
Bel LaPointe 2023-03-24 11:51:29 -06:00
parent e491cc5cbc
commit 6e1bfc177d
2 changed files with 7 additions and 2 deletions

View File

@ -6,6 +6,7 @@ import (
"io" "io"
"math/rand" "math/rand"
"os" "os"
"sort"
"time" "time"
"github.com/go-yaml/yaml" "github.com/go-yaml/yaml"
@ -87,11 +88,15 @@ func randomCharFromWeights(m map[byte]int) func() byte {
} }
sum += v sum += v
} }
sort.Slice(result, func(i, j int) bool {
return result[i].i < result[j].i
})
if sum <= 0 { if sum <= 0 {
panic("weights must total nonzero") panic("weights must total nonzero")
} }
return func() byte { return func() byte {
n := rand.Int() % sum r := rand.Int()
n := r % (sum + 1)
for _, v := range result { for _, v := range result {
n -= v.i n -= v.i
if n <= 0 { if n <= 0 {

View File

@ -8,7 +8,7 @@ import (
func TestRandomCharFromWeights(t *testing.T) { func TestRandomCharFromWeights(t *testing.T) {
weights := map[byte]int{ weights := map[byte]int{
'a': 1, 'a': 1,
'b': 99, 'b': 20,
} }
foo := randomCharFromWeights(weights) foo := randomCharFromWeights(weights)
for { for {