-
Notifications
You must be signed in to change notification settings - Fork 432
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Implement WeightedIndex, SliceRandom::choose_weighted and SliceRandom…
…::choose_weighted_mut
- Loading branch information
Showing
5 changed files
with
342 additions
and
28 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,182 @@ | ||
| // Copyright 2017 The Rust Project Developers. See the COPYRIGHT | ||
| // file at the top-level directory of this distribution and at | ||
| // https://rust-lang.org/COPYRIGHT. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or | ||
| // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license | ||
| // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your | ||
| // option. This file may not be copied, modified, or distributed | ||
| // except according to those terms. | ||
|
|
||
| use Rng; | ||
| use distributions::Distribution; | ||
| use distributions::uniform::{UniformSampler, SampleUniform, SampleBorrow}; | ||
| use ::core::cmp::PartialOrd; | ||
| use ::{Error, ErrorKind}; | ||
|
|
||
| // Note that this whole module is only imported if feature="alloc" is enabled. | ||
| #[cfg(not(feature="std"))] use alloc::Vec; | ||
|
|
||
| /// A distribution using weighted sampling to pick an discretely selected item. | ||
| /// | ||
| /// When a `WeightedIndex` is sampled from, it returns the index | ||
| /// of a random element from the iterator used when the `WeightedIndex` was | ||
| /// created. The chance of a given element being picked is proportional to the | ||
| /// value of the element. The weights can use any type `X` for which an | ||
| /// implementaiton of [`Uniform<X>`] exists. | ||
| /// | ||
| /// # Example | ||
| /// | ||
| /// ``` | ||
| /// use rand::prelude::*; | ||
| /// use rand::distributions::WeightedIndex; | ||
| /// | ||
| /// let choices = ['a', 'b', 'c']; | ||
| /// let weights = [2, 1, 1]; | ||
| /// let dist = WeightedIndex::new(&weights).unwrap(); | ||
| /// let mut rng = thread_rng(); | ||
| /// for _ in 0..100 { | ||
| /// // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c' | ||
| /// println!("{}", choices[dist.sample(&mut rng)]); | ||
| /// } | ||
| /// | ||
| /// let items = [('a', 0), ('b', 3), ('c', 7)]; | ||
| /// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap(); | ||
| /// for _ in 0..100 { | ||
| /// // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c' | ||
| /// println!("{}", items[dist2.sample(&mut rng)].0); | ||
| /// } | ||
| /// ``` | ||
| #[derive(Debug, Clone)] | ||
| pub struct WeightedIndex<X: SampleUniform + PartialOrd> { | ||
| cumulative_weights: Vec<X>, | ||
| weight_distribution: X::Sampler, | ||
| } | ||
|
|
||
| impl<X: SampleUniform + PartialOrd> WeightedIndex<X> { | ||
| /// Creates a new a `WeightedIndex` [`Distribution`] using the values | ||
| /// in `weights`. The weights can use any type `X` for which an | ||
| /// implementaiton of [`Uniform<X>`] exists. | ||
| /// | ||
| /// Returns an error if the iterator is empty, or its total value is 0. | ||
| /// | ||
| /// # Panics | ||
| /// | ||
| /// If a value in the iterator is `< 0`. | ||
| /// | ||
| /// [`Distribution`]: trait.Distribution.html | ||
| /// [`Uniform<X>`]: struct.Uniform.html | ||
| pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, Error> | ||
| where I: IntoIterator, | ||
| I::Item: SampleBorrow<X>, | ||
| X: for<'a> ::core::ops::AddAssign<&'a X> + | ||
| Clone + | ||
| Default { | ||
| let mut iter = weights.into_iter(); | ||
| let mut total_weight: X = iter.next() | ||
| .ok_or(Error::new(ErrorKind::Unexpected, "Empty iterator in WeightedIndex::new"))? | ||
| .borrow() | ||
| .clone(); | ||
|
|
||
| let zero = <X as Default>::default(); | ||
| let weights = iter.map(|w| { | ||
| assert!(*w.borrow() >= zero, "Negative weight in WeightedIndex::new"); | ||
| let prev_weight = total_weight.clone(); | ||
| total_weight += w.borrow(); | ||
| prev_weight | ||
| }).collect::<Vec<X>>(); | ||
|
|
||
| if total_weight == zero { | ||
| return Err(Error::new(ErrorKind::Unexpected, "Total weight is zero in WeightedIndex::new")); | ||
| } | ||
| let distr = X::Sampler::new(zero, total_weight); | ||
|
|
||
| Ok(WeightedIndex { cumulative_weights: weights, weight_distribution: distr }) | ||
| } | ||
| } | ||
|
|
||
| impl<X> Distribution<usize> for WeightedIndex<X> where | ||
| X: SampleUniform + PartialOrd { | ||
| fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize { | ||
| let chosen_weight = self.weight_distribution.sample(rng); | ||
| // Invariants: indexes in range [start, end] (inclusive) are candidate indexes | ||
| // cumulative_weights[start-1] <= chosen_weight | ||
| // chosen_weight < cumulative_weights[end] | ||
| // The returned index is the first one whose value is >= chosen_weight | ||
| let mut start = 0usize; | ||
| let mut end = self.cumulative_weights.len(); | ||
| while start < end { | ||
| let mid = (start + end) / 2; | ||
| if chosen_weight >= * unsafe { self.cumulative_weights.get_unchecked(mid) } { | ||
| start = mid + 1; | ||
| } else { | ||
| end = mid; | ||
| } | ||
| } | ||
| debug_assert_eq!(start, end); | ||
| start | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use super::*; | ||
| #[cfg(feature="std")] | ||
| use core::panic::catch_unwind; | ||
|
|
||
| #[test] | ||
| fn test_weightedindex() { | ||
| let mut r = ::test::rng(700); | ||
| const N_REPS: u32 = 5000; | ||
| let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7]; | ||
| let total_weight = weights.iter().sum::<u32>() as f32; | ||
|
|
||
| let verify = |result: [i32; 14]| { | ||
| for (i, count) in result.iter().enumerate() { | ||
| let exp = (weights[i] * N_REPS) as f32 / total_weight; | ||
| let mut err = (*count as f32 - exp).abs(); | ||
| if err != 0.0 { | ||
| err /= exp; | ||
| } | ||
| assert!(err <= 0.25); | ||
| } | ||
| }; | ||
|
|
||
| // WeightedIndex from vec | ||
| let mut chosen = [0i32; 14]; | ||
| let distr = WeightedIndex::new(weights.to_vec()).unwrap(); | ||
| for _ in 0..N_REPS { | ||
| chosen[distr.sample(&mut r)] += 1; | ||
| } | ||
| verify(chosen); | ||
|
|
||
| // WeightedIndex from slice | ||
| chosen = [0i32; 14]; | ||
| let distr = WeightedIndex::new(&weights[..]).unwrap(); | ||
| for _ in 0..N_REPS { | ||
| chosen[distr.sample(&mut r)] += 1; | ||
| } | ||
| verify(chosen); | ||
|
|
||
| // WeightedIndex from iterator | ||
| chosen = [0i32; 14]; | ||
| let distr = WeightedIndex::new(weights.iter()).unwrap(); | ||
| for _ in 0..N_REPS { | ||
| chosen[distr.sample(&mut r)] += 1; | ||
| } | ||
| verify(chosen); | ||
|
|
||
| assert!(WeightedIndex::new(&[10][0..0]).is_err()); | ||
| assert!(WeightedIndex::new(&[0]).is_err()); | ||
| } | ||
|
|
||
| #[test] | ||
| #[cfg(all(feature="std", | ||
| not(target_arch = "wasm32"), | ||
| not(target_arch = "asmjs")))] | ||
| fn test_weighted_assertions() { | ||
| assert!(catch_unwind(|| WeightedIndex::new(&[1, 2, 3])).is_ok()); | ||
| assert!(catch_unwind(|| WeightedIndex::new(&[10, -1, 10])).is_err()); | ||
| assert!(catch_unwind(|| WeightedIndex::new(&[1, -1])).is_err()); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.