steel_utils/random/
weighted_random.rs

1//! Weighted selection over an arbitrary slice.
2
3use crate::random::Random;
4
5/// Sums `weight` over `items`.
6///
7/// # Panics
8/// Panics if the sum exceeds `i32::MAX`.
9#[must_use]
10pub fn get_total_weight<T>(items: &[T], weight: impl Fn(&T) -> i32) -> i32 {
11    let total: i64 = items.iter().map(|item| i64::from(weight(item))).sum();
12    assert!(
13        total <= i64::from(i32::MAX),
14        "Sum of weights must be <= 2147483647"
15    );
16    total as i32
17}
18
19/// Picks an entry with probability proportional to its weight. A zero total
20/// weight returns `None` without consuming random values.
21///
22/// # Panics
23/// Panics if the total weight is negative.
24pub fn get_random_item<'a, T, R: Random>(
25    random: &mut R,
26    items: &'a [T],
27    weight: impl Fn(&T) -> i32,
28) -> Option<&'a T> {
29    let total_weight = get_total_weight(items, &weight);
30    assert!(total_weight >= 0, "Negative total weight in getRandomItem");
31    if total_weight == 0 {
32        return None;
33    }
34    let selection = random.next_i32_bounded(total_weight);
35    get_weighted_item(items, selection, weight)
36}
37
38/// Returns the entry whose cumulative weight range contains `index`.
39#[must_use]
40pub fn get_weighted_item<T>(items: &[T], mut index: i32, weight: impl Fn(&T) -> i32) -> Option<&T> {
41    items.iter().find(|item| {
42        index -= weight(item);
43        index < 0
44    })
45}
46
47#[cfg(test)]
48mod tests {
49    use crate::random::legacy_random::LegacyRandom;
50
51    use super::{get_random_item, get_total_weight, get_weighted_item};
52
53    const ITEMS: [(&str, i32); 3] = [("common", 3), ("rare", 1), ("never", 0)];
54
55    fn weight(item: &(&str, i32)) -> i32 {
56        item.1
57    }
58
59    #[test]
60    fn total_weight_sums_every_entry() {
61        assert_eq!(get_total_weight(&ITEMS, weight), 4);
62        assert_eq!(get_total_weight::<(&str, i32)>(&[], weight), 0);
63    }
64
65    #[test]
66    fn weighted_item_walks_cumulative_weights() {
67        for index in 0..3 {
68            assert_eq!(get_weighted_item(&ITEMS, index, weight), Some(&ITEMS[0]));
69        }
70        assert_eq!(get_weighted_item(&ITEMS, 3, weight), Some(&ITEMS[1]));
71        assert_eq!(get_weighted_item(&ITEMS, 4, weight), None);
72    }
73
74    #[test]
75    fn zero_total_weight_yields_nothing_without_consuming_the_random() {
76        let mut random = LegacyRandom::from_seed(5);
77        let before = random.get_seed();
78        let empty: [(&str, i32); 1] = [("never", 0)];
79        assert_eq!(get_random_item(&mut random, &empty, weight), None);
80        assert_eq!(random.get_seed(), before);
81    }
82
83    #[test]
84    fn random_item_is_deterministic_for_a_seed() {
85        let first = get_random_item(&mut LegacyRandom::from_seed(1234), &ITEMS, weight);
86        let second = get_random_item(&mut LegacyRandom::from_seed(1234), &ITEMS, weight);
87        assert!(first.is_some());
88        assert_eq!(first, second);
89    }
90}
steel_utils/random/weighted_random.rs

steel_utils/random/
weighted_random.rs
