weightedrand

README

weightedrand ⚖

Fast weighted random selection for Go.

Randomly selects an element from some kind of list, where the chances of each element to be selected are not equal, but rather defined by relative "weights" (or probabilities). This is called weighted random selection.

Usage

import (
    /* ...snip... */
    "github.com/mroth/weightedrand/v2"
)

func main() {
    chooser, _ := weightedrand.NewChooser(
        weightedrand.NewChoice('🍒', 0),
        weightedrand.NewChoice('🍋', 1),
        weightedrand.NewChoice('🍊', 1),
        weightedrand.NewChoice('🍉', 3),
        weightedrand.NewChoice('🥑', 5),
    )
    // The following will print 🍋 and 🍊 with 0.1 probability, 🍉 with 0.3
    // probability, and 🥑 with 0.5 probability. 🍒 will never be printed. (Note
    // the weights don't have to add up to 10, that was just done here to make
    // the example easier to read.)
    result := chooser.Pick()
    fmt.Println(result)
}

Performance

The existing Go library that has a comparable implementation of this is github.com/jmcvetta/randutil, which optimizes for the single operation case. In contrast, this library creates a presorted cache optimized for binary search, allowing repeated selections from the same set to be significantly faster, especially for large data sets.

Comparison of this library versus randutil.ChooseWeighted on my workstation. For repeated samplings from large collections, weightedrand will be much quicker:

Num choices	randutil	weightedrand	weightedrand -cpu=8*
10	201 ns/op	38 ns/op	2.9 ns/op
100	267 ns/op	51 ns/op	4.1 ns/op
1,000	1012 ns/op	67 ns/op	5.4 ns/op
10,000	8683 ns/op	83 ns/op	6.9 ns/op
100,000	123500 ns/op	105 ns/op	12.0 ns/op
1,000,000	2399614 ns/op	218 ns/op	17.2 ns/op
10,000,000	26804440 ns/op	432 ns/op	35.1 ns/op

*: Since v0.3.0 weightedrand can efficiently utilize a single Chooser across multiple CPU cores in parallel, making it even faster in overall throughput. See PR#2 for details. Informal benchmarks conducted on an Intel Xeon W-2140B CPU (8 core @ 3.2GHz, hyperthreading enabled).

Don't be mislead by these numbers into thinking weightedrand is always the right choice! If you are only picking from the same distribution once, randutil will be faster. weightedrand optimizes for repeated calls at the expense of some initialization time and memory storage.

Requirements

weightedrand >= v2 requires go1.18 or greater. For support on earlier versions of go, use weightedrand v1.

Credits

To better understand the algorithm used in this library (as well as the one used in randutil) check out this great blog post: Weighted random generation in Python.

Documentation

Overview

Package weightedrand contains a performant data structure and algorithm used to randomly select an element from some kind of list, where the chances of each element to be selected not being equal, but defined by relative "weights" (or probabilities). This is called weighted random selection.

This package creates a presorted cache optimized for binary search, allowing for repeated selections from the same set to be significantly faster, especially for large data sets.

Example

In this example, we create a Chooser to pick from amongst various emoji fruit runes. We assign a numeric weight to each choice. These weights are relative, not on any absolute scoring system. In this trivial case, we will assign a weight of 0 to all but one fruit, so that the output will be predictable.

chooser, _ := NewChooser(
	NewChoice('🍋', 0),
	NewChoice('🍊', 0),
	NewChoice('🍉', 0),
	NewChoice('🥑', 42),
)
fruit := chooser.Pick()
fmt.Printf("%c", fruit)

Output:

🥑

Index

• type Choice
  • func NewChoice[T any, W integer](item T, weight W) Choice[T, W]
• type Chooser
  • func NewChooser[T any, W integer](choices ...Choice[T, W]) (*Chooser[T, W], error)
  • func (c Chooser[T, W]) Pick() T
  • func (c Chooser[T, W]) PickSource(rs *rand.Rand) T

Examples

• Package

Types

type Choice

type Choice[T any, W integer] struct {
	Item   T
	Weight W
}

Choice is a generic wrapper that can be used to add weights for any item.

func NewChoice

func NewChoice[T any, W integer](item T, weight W) Choice[T, W]

NewChoice creates a new Choice with specified item and weight.

type Chooser

type Chooser[T any, W integer] struct {
	// contains filtered or unexported fields
}

A Chooser caches many possible Choices in a structure designed to improve performance on repeated calls for weighted random selection.

func NewChooser

func NewChooser[T any, W integer](choices ...Choice[T, W]) (*Chooser[T, W], error)

NewChooser initializes a new Chooser for picking from the provided choices.

func (Chooser[T, W]) Pick

func (c Chooser[T, W]) Pick() T

Pick returns a single weighted random Choice.Item from the Chooser.

Utilizes global rand as the source of randomness. Safe for concurrent usage.

func (Chooser[T, W]) PickSource

func (c Chooser[T, W]) PickSource(rs *rand.Rand) T

PickSource returns a single weighted random Choice.Item from the Chooser, utilizing the provided *rand.Rand source rs for randomness.

The primary use-case for this is avoid lock contention from the global random source if utilizing Chooser(s) from multiple goroutines in extremely high-throughput situations.

It is the responsibility of the caller to ensure the provided rand.Source is free from thread safety issues.

Deprecated: Since go1.21 global rand no longer suffers from lock contention when used in multiple high throughput goroutines, as long as you don't manually seed it. Use Chooser.Pick instead.