wallrnd/
chooser.rs

1use rand::{rngs::ThreadRng, Rng};
2
3#[derive(Clone)]
4pub struct Chooser<T: Clone>(usize, Vec<(T, usize)>);
5
6impl<T: Clone> Chooser<T> {
7    /// Empty Chooser
8    pub fn default() -> Self {
9        Self(0, Vec::new())
10    }
11
12    /// Create Chooser from weighted items
13    pub fn new(mut v: Vec<(T, usize)>) -> Self {
14        let mut sum = 0;
15        for (_, w) in &mut v {
16            sum += *w;
17            *w = sum;
18        }
19        if !v.is_empty() {
20            Self(sum, v)
21        } else {
22            Self::default()
23        }
24    }
25
26    /// Pick a random item (weighted)
27    pub fn choose(&self, rng: &mut ThreadRng) -> Option<T> {
28        if self.1.is_empty() {
29            None
30        } else {
31            let choice = rng.gen_range(0, self.0);
32            Some(self.dichotomy(choice, 0, self.1.len()))
33        }
34    }
35
36    fn dichotomy(&self, target: usize, inf: usize, sup: usize) -> T {
37        if inf == sup {
38            self.1[inf].0.clone()
39        } else if inf + 1 == sup {
40            if self.1[inf].1 < target {
41                self.1[inf + 1].0.clone()
42            } else {
43                self.1[inf].0.clone()
44            }
45        } else {
46            let mid = (sup + inf) / 2;
47            if self.1[mid].1 > target {
48                self.dichotomy(target, inf, mid)
49            } else {
50                self.dichotomy(target, mid, sup)
51            }
52        }
53    }
54
55    /// Get items with their weights as a copy
56    pub fn extract(&self) -> Vec<(T, usize)> {
57        let mut cpy = self.1.clone();
58        let n = cpy.len();
59        for i in 1..n {
60            cpy[n - i].1 -= cpy[n - i - 1].1;
61        }
62        cpy
63    }
64
65    /// Add new item
66    pub fn push(&mut self, item: T, w: usize) {
67        if w == 0 {
68            panic!("In call to Chooser::push, 0 is not a valid weight");
69        }
70        self.0 += w;
71        self.1.push((item, self.0));
72    }
73
74    /// Add vector of new items
75    pub fn append(&mut self, items: Vec<(T, usize)>) {
76        for (item, w) in items {
77            self.push(item, w);
78        }
79    }
80}
wallrnd/chooser.rs

wallrnd/
chooser.rs
