weighted_random_list/

lib.rs

//! `WeightedRandomList` enables you to to randomly select elements from
//! your collection while
//!
//! # Example select mirror
//!
//! In this example we would like to give more traffic/weight to the the system
//! "mirror-fast", less to the slow mirrors and the least to the main archive:
//!
//! ```rust
//! extern crate weighted_random_list;
//!
//! use weighted_random_list::WeightedRandomList;
//!
//! fn main() {
//!     let list = [
//!         (1, "https://source.example.net/archive"),
//!         (10, "https://mirror-slow0.example.net/archive"),
//!         (10, "https://mirror-slow1.example.net/archive"),
//!         (100, "https://mirror-fast.example.net/archive"),
//!     ];
//!
//!     let mirrors = list.iter()
//!             .map(|(weight, url)| (*weight, url.to_string()))
//!             .collect::<WeightedRandomList<String>>();
//!
//!
//!     let random_choice = mirrors.get_random();
//!     println!("Using {:?} this time", random_choice);
//! }
//! ```

use std::iter::FromIterator;

use rand::Rng;

/// Stores the data and keeps track of the chances.
/// It is built so that you can add and remove elements withouth the need of
/// recalculation.
pub struct WeightedRandomList<T> {
    sum_weights: usize,
    data: Vec<(usize, T)>,
}

impl<T> WeightedRandomList<T> {
    pub fn new() -> WeightedRandomList<T> {
        WeightedRandomList {
            sum_weights: 0,
            data: Vec::new(),
        }
    }

    /// Add an element to the collection with a given weight
    ///
    /// Panics if the sum of the weights exceed the holding datatype
    pub fn push(&mut self, weight: usize, new: T) {
        self.sum_weights = self.sum_weights.checked_add(weight)
                                .expect("WeightedRandomList.push() sum of weights exceided data type");
        self.data.push((weight , new))
    }

    /// Get a random value from the collection or `None` if it is empty.
    ///
    /// Accsess time is `O(n)`, however the linear storage compensates for
    /// this overhead.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use weighted_random_list::WeightedRandomList;
    /// let mut l = WeightedRandomList::new();
    /// assert_eq!(None, l.get_random());
    ///
    /// l.push(100, 42usize);
    /// assert_eq!(Some(&42), l.get_random());
    /// ```
    pub fn get_random(&self) -> Option<&T> {
        if self.data.is_empty() {
            return None;
        }
        let r = rand::thread_rng().gen_range(0, self.sum_weights);

        let mut local_sum = 0;
        for (weight, data) in &self.data {
            local_sum += weight;
            if local_sum >= r {
                return Some(data)
            }
        }
        None
    }
}

/// If the stored type `T` supports `PartialEq` and `Clone` we can manipulate
/// the contents of this container.
/// Therefore shifting the chances of one element to be selected.
///
/// # Example adjust weight after insertion
///
/// ```rust
/// # use weighted_random_list::WeightedRandomList;
/// let mut l = WeightedRandomList::new();
/// l.push(10, 1);
/// l.push(10, 2);
/// l.push(100, 3);
/// // total weights is now 120
///
/// {
///     let mut three = l.entry_first(&3);
///     three.set_weight(80);
/// }
/// // total weights is now 100
/// ```
impl<T> WeightedRandomList<T> where T: std::cmp::PartialEq + Clone {

    /// Find the first entry in the list matching `needle` and manipulate
    /// potential entries:
    ///
    /// ```rust
    /// # use weighted_random_list::WeightedRandomList;
    /// let mut l = WeightedRandomList::new();
    ///
    /// l.entry_first(&42).or_insert_with_weight(10);
    /// ```
    pub fn entry_first(&mut self, needle: &T) -> Entry<T> {
        for (index, (_weight, data)) in self.data.iter().enumerate() {
            if needle == data {
                return Entry::Occupied {
                    arena: self,
                    index,
                }
            }
        }

        Entry::Vacant {
            arena: self,
            needle: needle.clone(),
        }
    }
}

/// Update or insert one entry inside the the list.
pub enum Entry<'container, T> {
    Vacant {
        arena: &'container mut WeightedRandomList<T>,
        needle: T,
    },
    Occupied {
        arena: &'container mut WeightedRandomList<T>,
        index: usize,
    }
}

impl<'container, T> Entry<'container, T> {
    /// Panics if `Entry` is `Vacant`
    /// ```rust
    /// # use weighted_random_list::WeightedRandomList;
    /// let mut l = WeightedRandomList::new();
    ///
    /// l.push(12, 42);
    /// let mut element = l.entry_first(&42).or_insert_with_weight(10);
    ///
    /// element.set_weight(100);
    /// ```
    pub fn set_weight(&mut self, new_weight: usize) {
        use Entry::*;
        match self {
            Vacant { .. } => { unimplemented!("set_weight on variant Vacant"); }
            Occupied { arena, index } => {
                let (weight, _data) = &mut arena.data
                    //[*index];
                    .get_mut(*index)
                    .expect("unable to directly access data from Entry view");
                arena.sum_weights -= *weight;
                arena.sum_weights += new_weight;
                *weight = new_weight;
            }
        }
    }

    /// Inserts if `Entry` is `Vacant`.
    /// Does nothing if found inside the list.

    /// ```rust
    /// # use weighted_random_list::WeightedRandomList;
    /// let mut l = WeightedRandomList::new();
    ///
    /// // add 42 with weight 10
    /// l.entry_first(&42).or_insert_with_weight(10);
    /// ```
    pub fn or_insert_with_weight(self, weight: usize) -> Self {
        use Entry::*;
        match self {
            Occupied { .. } => { self }
            Vacant { arena, needle } => {
                let index = arena.data.len();
                arena.push(weight, needle);

                Occupied {
                    arena,
                    index,
                }
            }
        }
    }

    /// Remove `Entry` if it exists.
    ///
    /// Errors if it does not exist.
    ///
    /// ```rust
    /// # use weighted_random_list::WeightedRandomList;
    /// let mut l = WeightedRandomList::new();
    /// l.push(12, 11);
    /// l.push(74, 22);
    /// l.push(25, 33);
    ///
    /// {
    ///     let second = l.entry_first(&22);
    ///     assert!(second.delete().is_ok(), "element was found and deleted");
    /// }
    ///
    /// {
    ///     let second = l.entry_first(&22);
    ///     assert!(second.delete().is_err(), "element alread gone");
    /// }
    /// ```
    pub fn delete(self) -> Result<(), ()> {
        use Entry::*;
        match self {
            Vacant { .. } => { Err(()) }
            Occupied { arena, index } => {
                let len = arena.data.len();
                if len != (index + 1) {
                    arena.data.swap(index, len - 1);
                }
                let (weight, _data) = arena.data.pop().expect("must be the last element");
                arena.sum_weights -= weight;
                Ok(())
            }
        }
    }
}

/// Enable `clollect`
impl<T> FromIterator<(usize, T)> for WeightedRandomList<T> {
    fn from_iter<I: IntoIterator<Item=(usize, T)>>(iter: I) -> Self {
        let mut c = WeightedRandomList::new();

        for (weight, data) in iter {
            c.push(weight, data);
        }

        c
    }
}
/* TODO implement in a way to be able to collect &str from static list
impl<'a, T: 'a> FromIterator<&'a(usize, &T)> for WeightedRandomList<T> {
    fn from_iter<I: IntoIterator<Item=&'a (usize, &T)>>(iter: I) -> Self {
        let mut c = WeightedRandomList::new();

        for (weight, data) in iter {
            c.push(*weight, *data);
        }

        c
    }
}
*/

#[cfg(test)]
mod tests {
    use crate::*;
    use std::collections::HashMap;

    #[test]
    fn simple_inserts() {
        let mut l = WeightedRandomList::new();
        assert_eq!(None, l.get_random());

        l.push(100, 42usize);
        assert_eq!(Some(&42), l.get_random());

        l.push(0, 23);
        for _ in 0..1000 {
            assert_eq!(Some(&42), l.get_random());
        }

        l.push(900, 12);
        let results = (0..1000).map(|_| l.get_random().unwrap().clone())
                .fold(HashMap::new(), |mut map, n| {
                    let count = map.entry(n).or_insert(0);
                    *count += 1;
                    map
                });
        assert!(results[&42] < results[&12]);
    }

    #[test]
    fn modify_weight() {
        let mut l = WeightedRandomList::new();
        l.push(10, 1);
        l.push(10, 2);
        l.push(100, 3);
        assert_eq!(120, l.sum_weights);

        {
            let mut three = l.entry_first(&3);
            three.set_weight(80);
        }
        assert_eq!(100, l.sum_weights);

        let results = (0..1000).map(|_| l.get_random().unwrap().clone())
                .fold(HashMap::new(), |mut map, n| {
                    let count = map.entry(n).or_insert(0);
                    *count += 1;
                    map
                });
        assert!(results[&1] + results[&2] < results[&3]);
    }

    #[test]
    fn or_insert() {
        let mut l = WeightedRandomList::new();

        l.entry_first(&1).or_insert_with_weight(10);

        {
            assert_eq!(10, l.sum_weights);
            let mut two = l.entry_first(&2).or_insert_with_weight(1);
            two.set_weight(10);
            assert_eq!(20, l.sum_weights);
        }

        l.entry_first(&3).or_insert_with_weight(80);

        let results = (0..1000).map(|_| l.get_random().unwrap().clone())
                .fold(HashMap::new(), |mut map, n| {
                    let count = map.entry(n).or_insert(0);
                    *count += 1;
                    map
                });
        assert!(results[&1] + results[&2] < results[&3]);
    }

    #[test]
    fn delete_last() {
        let mut l = WeightedRandomList::new();
        l.push(100, 3);
        assert_eq!(1, l.data.len());
        assert_eq!(100, l.sum_weights);

        {
            let three = l.entry_first(&3);
            assert!(three.delete().is_ok());
        }
        assert_eq!(0, l.sum_weights);
        assert_eq!(0, l.data.len());
    }
    #[test]
    fn delete_mid() {
        let mut l = WeightedRandomList::new();
        assert_eq!(0, l.data.len());
        l.push(10, 1);
        assert_eq!(1, l.data.len());
        l.push(10, 2);
        assert_eq!(2, l.data.len());
        l.push(100, 3);
        assert_eq!(3, l.data.len());
        assert_eq!(120, l.sum_weights);

        {
            let two = l.entry_first(&2);
            assert!(two.delete().is_ok());
        }
        assert_eq!(110, l.sum_weights); assert_eq!(2, l.data.len());
    }
    #[test]
    fn delete_end() {
        let mut l = WeightedRandomList::new();
        assert_eq!(0, l.data.len());
        l.push(10, 1);
        assert_eq!(1, l.data.len());
        l.push(10, 2);
        assert_eq!(2, l.data.len());
        l.push(100, 3);
        assert_eq!(3, l.data.len());
        assert_eq!(120, l.sum_weights);

        {
            let three = l.entry_first(&3);
            assert!(three.delete().is_ok());
        }
        assert_eq!(20, l.sum_weights);
        assert_eq!(2, l.data.len());
    }
    #[test]
    fn delete_not_found() {
        let mut l = WeightedRandomList::new();
        assert_eq!(0, l.data.len());
        l.push(10, 1);
        assert_eq!(1, l.data.len());
        l.push(10, 2);
        assert_eq!(2, l.data.len());
        l.push(100, 3);
        assert_eq!(3, l.data.len());
        assert_eq!(120, l.sum_weights);

        {
            let absent = l.entry_first(&22);
            assert!(absent.delete().is_err());
        }
        assert_eq!(120, l.sum_weights);
        assert_eq!(3, l.data.len());
    }

    #[test]
    fn collect() {
        // make sure this compiles
        let list = [
            (1, "https://source.example.net/archive"),
            (10, "https://mirror-slow0.example.net/archive"),
            (10, "https://mirror-slow1.example.net/archive"),
            (100, "https://mirror-fast.example.net/archive"),
        ];

        let collected = list.iter()
                .map(|(weight, url)| (*weight, url.to_string()))
                .collect::<WeightedRandomList<String>>();

        assert_eq!(list.len(), collected.data.len());
        for ((w, url), (wc, urlc)) in list.iter().zip(collected.data.iter()) {
            assert_eq!(w, wc);
            assert_eq!(url, urlc);
        }
        assert_eq!(list.iter().map(|(w,_)| w).sum::<usize>(), collected.sum_weights);
    }
}