File tree 1 file changed +24
-3
lines changed
1 file changed +24
-3
lines changed Original file line number Diff line number Diff line change @@ -123,10 +123,24 @@ pub trait IteratorExt: Iterator + Sized {
123
123
/// Complexity is `O(n)`, where `n` is the length of the iterator.
124
124
/// This likely consumes multiple random numbers, but the exact number
125
125
/// 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 >
127
127
where R : Rng + ?Sized
128
128
{
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
+ }
130
144
}
131
145
132
146
/// Collects `amount` values at random from the iterator into a supplied
@@ -485,7 +499,14 @@ mod test {
485
499
#[ test]
486
500
fn test_choose ( ) {
487
501
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 ) ) ;
489
510
490
511
let v: & [ isize ] = & [ ] ;
491
512
assert_eq ! ( v. choose( & mut r) , None ) ;
You can’t perform that action at this time.
0 commit comments