weighted_die.rs

use rand::Rng;

struct WeightedDie {
    n: usize,
    prob: Vec<f64>,
    alias: Vec<usize>,
}

impl WeightedDie {
    fn new(weights: &Vec<i32>) -> Self {
        let n = weights.len();
        let sum = weights.iter().sum::<i32>() as f64;
        let mut p: Vec<f64> = weights
            .iter()
            .map(|w| (*w as f64) / sum * (n as f64))
            .collect();

        let mut prob: Vec<f64> = vec![0_f64; weights.len()];
        let mut alias: Vec<usize> = vec![0; weights.len()];

        let mut small: Vec<usize> = Vec::new();
        let mut large: Vec<usize> = Vec::new();
        for (i, pp) in p.iter().enumerate() {
            if *pp < 1_f64 {
                small.push(i);
            } else {
                large.push(i);
            }
        }

        while !small.is_empty() && !large.is_empty() {
            let l = small.pop().unwrap();
            let g = large.pop().unwrap();
            prob[l] = p[l];
            alias[l] = g;
            p[g] = p[g] + p[l] - 1_f64;
            if p[g] < 1_f64 {
                small.push(g);
            } else {
                large.push(g);
            }
        }
        while let Some(g) = large.pop() {
            prob[g] = 1_f64;
        }
        while let Some(l) = small.pop() {
            prob[l] = 1_f64;
        }

        Self { n, prob, alias }
    }

    fn roll(&self) -> usize {
        let i = rand::rng().random_range(..self.n);
        let r = rand::rng().random_range(0_f64..1_f64);
        if r < self.prob[i] { i } else { self.alias[i] }
    }
}

fn main() {
    let weights = vec![1, 2, 3, 10, 30, 3, 5, 5, 10, 10];
    let die = WeightedDie::new(&weights);
    let mut hits = vec![0; weights.len()];
    for _ in 0..100 {
        let rolled = die.roll();
        hits[rolled] += 1;
    }

    for (i, n) in hits.iter().enumerate() {
        println!("{} (weight: {}) -> {} hits", i, weights[i], n);
    }
}

weighted_die.ts

export class WeightedDie {
  private readonly n: number
  private readonly prob: readonly number[]
  private readonly alias: readonly number[]

  public static fromWeightFn(n: number, weight: (i: number) => number): WeightedDie {
    return new WeightedDie(new Array(n).fill(0).map((_, i) => weight(i)))
  }

  public constructor(weights: number[]) {
    // https://en.wikipedia.org/wiki/Alias_method
    // https://stackoverflow.com/a/39199014
    const n = weights.length
    const sum = weights.reduce((acc, cur) => acc + cur, 0)
    const p = weights.map((w) => w / sum * n)

    const prob = new Array(n).fill(0)
    const alias = new Array(n).fill(0)

    const small: number[] = []
    const large: number[] = []
    p.forEach((pp, i) => {
      if (pp < 1) small.push(i)
      else large.push(i)
    })

    while (small.length > 0 && large.length > 0) {
      const l = small.shift()!
      const g = large.shift()!
      prob[l] = p[l]
      alias[l] = g
      p[g] = (p[g] + p[l]) - 1
      if (p[g] < 1) small.push(g)
      else large.push(g)
    }
    while (large.length > 0) {
      const g = large.shift()!
      prob[g] = 1
    }
    while (small.length > 0) {
      const l = small.shift()!
      prob[l] = 1
    }

    this.n = n
    this.prob = prob
    this.alias = alias
  }

  public roll(): number {
    const i = randomIntn(this.n)
    if (Math.random() < this.prob[i]) return i
    else return this.alias[i]
  }
}