Weighted Random Choice

package main

import (
	"errors"
	"fmt"
	"math/rand"
)

type Choice[T any] struct {
	Item   T
	Weight float64
}

func Pick[T any](r *rand.Rand, choices []Choice[T]) (T, error) {
	var zero T
	if len(choices) == 0 {
		return zero, errors.New("no choices")
	}
	total := 0.0
	for _, choice := range choices {
		if choice.Weight < 0 {
			return zero, fmt.Errorf("negative weight for %v", choice.Item)
		}
		total += choice.Weight
	}
	if total <= 0 {
		return zero, errors.New("total weight must be positive")
	}
	target := r.Float64() * total
	for _, choice := range choices {
		if target < choice.Weight {
			return choice.Item, nil
		}
		target -= choice.Weight
	}
	return choices[len(choices)-1].Item, nil
}

func main() {
	rng := rand.New(rand.NewSource(42))
	choices := []Choice[string]{
		{Item: "canary", Weight: 1},
		{Item: "stable", Weight: 3},
		{Item: "legacy", Weight: 6},
	}
	counts := map[string]int{}
	for i := 0; i < 10000; i++ {
		item, err := Pick(rng, choices)
		if err != nil {
			panic(err)
		}
		counts[item]++
	}
	for _, choice := range choices {
		fmt.Printf("%s weight=%.0f picked=%d (%.1f%%)\n", choice.Item, choice.Weight, counts[choice.Item], float64(counts[choice.Item])/100)
	}
	_, err := Pick(rng, []Choice[string]{{Item: "bad", Weight: -1}})
	fmt.Println("negative weight rejected:", err != nil)
}