Skip to content

Commit

Permalink
Implement IteratorExt::choose
Browse files Browse the repository at this point in the history
  • Loading branch information
dhardy committed May 27, 2018
1 parent 922bb09 commit d697fd0
Showing 1 changed file with 24 additions and 3 deletions.
27 changes: 24 additions & 3 deletions src/seq.rs
Expand Up @@ -123,10 +123,24 @@ pub trait IteratorExt: Iterator + Sized {
/// Complexity is `O(n)`, where `n` is the length of the iterator.
/// This likely consumes multiple random numbers, but the exact number
/// is unspecified.
fn choose<R>(self, rng: &mut R) -> Option<Self::Item>
fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
where R: Rng + ?Sized
{
unimplemented!()
if let Some(elem) = self.next() {
let mut result = elem;

// Continue until the iterator is exhausted
for (i, elem) in self.enumerate() {
let k = rng.gen_range(0, i + 2);
if k == 0 {
result = elem;
}
}

Some(result)
} else {
None
}
}

/// Collects `amount` values at random from the iterator into a supplied
Expand Down Expand Up @@ -485,7 +499,14 @@ mod test {
#[test]
fn test_choose() {
let mut r = ::test::rng(107);
assert_eq!([1, 1, 1].choose(&mut r).map(|&x|x), Some(1));
assert_eq!([1, 1, 1].choose(&mut r), Some(&1));

let mut v = [2];
v.choose_mut(&mut r).map(|x| *x = 5);
assert_eq!(v[0], 5);

let v = [3, 3, 3, 3];
assert_eq!(v.iter().choose(&mut r), Some(&3));

let v: &[isize] = &[];
assert_eq!(v.choose(&mut r), None);
Expand Down

0 comments on commit d697fd0

Please sign in to comment.