weighted_selector/
lib.rs

1/*
2 * Copyright (c) Peter Bjorklund. All rights reserved. https://github.com/piot/weighted-selector
3 * Licensed under the MIT License. See LICENSE in the project root for license information.
4 */
5
6pub struct WeightedSelectorIndex {
7    weights: Vec<usize>,
8    total_size: usize,
9}
10
11impl WeightedSelectorIndex {
12    pub fn new(weights: Vec<usize>) -> Self {
13        let total_size = weights.iter().sum();
14        Self {
15            weights,
16            total_size,
17        }
18    }
19
20    pub fn total(&self) -> usize {
21        self.total_size
22    }
23
24    pub fn select(&self, value: usize) -> Option<usize> {
25        let mut cumulative = 0;
26
27        for (i, &weight) in self.weights.iter().enumerate() {
28            cumulative += weight;
29            if value < cumulative {
30                return Some(i);
31            }
32        }
33
34        None
35    }
36}
37
38pub struct WeightedSelector<T> {
39    enums: Vec<(usize, T)>,
40    total_size: usize,
41}
42
43impl<T> WeightedSelector<T> {
44    pub fn new(enums: Vec<(usize, T)>) -> Self {
45        let total_size = enums.iter().map(|(weight, _)| weight).sum();
46        Self { enums, total_size }
47    }
48
49    #[inline]
50    pub fn total(&self) -> usize {
51        self.total_size
52    }
53
54    #[inline]
55    pub fn select(&self, value: usize) -> Option<&T> {
56        let mut cumulative = 0;
57
58        for (weight, enum_value) in &self.enums {
59            cumulative += weight;
60            if value < cumulative {
61                return Some(enum_value);
62            }
63        }
64
65        None
66    }
67}
weighted_selector/lib.rs

weighted_selector/
lib.rs
