weighted_selector/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
/*
 * Copyright (c) Peter Bjorklund. All rights reserved. https://github.com/piot/weighted-selector
 * Licensed under the MIT License. See LICENSE in the project root for license information.
 */

pub struct WeightedSelectorIndex {
    weights: Vec<usize>,
}

impl WeightedSelectorIndex {
    pub const fn new(weights: Vec<usize>) -> Self {
        Self { weights }
    }

    pub fn select(&self, value: usize) -> Option<usize> {
        let mut cumulative = 0;

        for (i, &weight) in self.weights.iter().enumerate() {
            cumulative += weight;
            if value < cumulative {
                return Some(i);
            }
        }

        None
    }
}

pub struct WeightedSelector<T> {
    enums: Vec<(usize, T)>,
}

impl<T> WeightedSelector<T> {
    pub const fn new(enums: Vec<(usize, T)>) -> Self {
        Self { enums }
    }

    pub fn select(&self, value: usize) -> Option<&T> {
        let mut cumulative = 0;

        for (weight, enum_value) in &self.enums {
            cumulative += weight;
            if value < cumulative {
                return Some(enum_value);
            }
        }

        None
    }
}