Skip to content

Commit d697fd0

Browse files
committed
Implement IteratorExt::choose
1 parent 922bb09 commit d697fd0

File tree

1 file changed

+24
-3
lines changed

1 file changed

+24
-3
lines changed

src/seq.rs

+24-3
Original file line numberDiff line numberDiff line change
@@ -123,10 +123,24 @@ pub trait IteratorExt: Iterator + Sized {
123123
/// Complexity is `O(n)`, where `n` is the length of the iterator.
124124
/// This likely consumes multiple random numbers, but the exact number
125125
/// is unspecified.
126-
fn choose<R>(self, rng: &mut R) -> Option<Self::Item>
126+
fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
127127
where R: Rng + ?Sized
128128
{
129-
unimplemented!()
129+
if let Some(elem) = self.next() {
130+
let mut result = elem;
131+
132+
// Continue until the iterator is exhausted
133+
for (i, elem) in self.enumerate() {
134+
let k = rng.gen_range(0, i + 2);
135+
if k == 0 {
136+
result = elem;
137+
}
138+
}
139+
140+
Some(result)
141+
} else {
142+
None
143+
}
130144
}
131145

132146
/// Collects `amount` values at random from the iterator into a supplied
@@ -485,7 +499,14 @@ mod test {
485499
#[test]
486500
fn test_choose() {
487501
let mut r = ::test::rng(107);
488-
assert_eq!([1, 1, 1].choose(&mut r).map(|&x|x), Some(1));
502+
assert_eq!([1, 1, 1].choose(&mut r), Some(&1));
503+
504+
let mut v = [2];
505+
v.choose_mut(&mut r).map(|x| *x = 5);
506+
assert_eq!(v[0], 5);
507+
508+
let v = [3, 3, 3, 3];
509+
assert_eq!(v.iter().choose(&mut r), Some(&3));
489510

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

0 commit comments

Comments
 (0)