use rand::{Rng, RngExt};
pub fn sample<R: Rng + ?Sized>(rng: &mut R, probs: impl ExactSizeIterator<Item = f32>) -> usize {
let last = probs.len() - 1;
let mut remaining: f32 = rng.random();
for (index, prob) in probs.take(last).enumerate() {
if prob < remaining {
remaining -= prob;
} else {
return index;
}
}
last
}
#[cfg(test)]
mod tests {
use super::sample;
use rand::rng;
#[test]
fn test_invalid_small() {
for _ in 0..1000 {
assert!((0..3).contains(&sample(&mut rng(), [0.2, 0.3, 0.4].into_iter())));
}
}
#[test]
fn test_invalid_large() {
for _ in 0..1000 {
assert!((0..3).contains(&sample(&mut rng(), [0.2, 0.7, 0.4].into_iter())));
}
}
}