]>
Commit | Line | Data |
---|---|---|
0731742a XL |
1 | // Copyright 2018 Developers of the Rand project. |
2 | // | |
3 | // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or | |
4 | // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license | |
5 | // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your | |
6 | // option. This file may not be copied, modified, or distributed | |
7 | // except according to those terms. | |
8 | ||
9 | //! Functions for randomly accessing and sampling sequences. | |
10 | //! | |
11 | //! TODO: module doc | |
12 | ||
13 | ||
14 | #[cfg(feature="alloc")] pub mod index; | |
15 | ||
16 | #[cfg(feature="alloc")] use core::ops::Index; | |
17 | ||
18 | #[cfg(all(feature="alloc", not(feature="std")))] use alloc::vec::Vec; | |
19 | ||
20 | use Rng; | |
21 | #[cfg(feature="alloc")] use distributions::WeightedError; | |
22 | #[cfg(feature="alloc")] use distributions::uniform::{SampleUniform, SampleBorrow}; | |
23 | ||
24 | /// Extension trait on slices, providing random mutation and sampling methods. | |
25 | /// | |
26 | /// An implementation is provided for slices. This may also be implementable for | |
27 | /// other types. | |
28 | pub trait SliceRandom { | |
29 | /// The element type. | |
30 | type Item; | |
31 | ||
32 | /// Returns a reference to one random element of the slice, or `None` if the | |
33 | /// slice is empty. | |
34 | /// | |
35 | /// Depending on the implementation, complexity is expected to be `O(1)`. | |
36 | /// | |
37 | /// # Example | |
38 | /// | |
39 | /// ``` | |
40 | /// use rand::thread_rng; | |
41 | /// use rand::seq::SliceRandom; | |
42 | /// | |
43 | /// let choices = [1, 2, 4, 8, 16, 32]; | |
44 | /// let mut rng = thread_rng(); | |
45 | /// println!("{:?}", choices.choose(&mut rng)); | |
46 | /// assert_eq!(choices[..0].choose(&mut rng), None); | |
47 | /// ``` | |
48 | fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item> | |
49 | where R: Rng + ?Sized; | |
50 | ||
51 | /// Returns a mutable reference to one random element of the slice, or | |
52 | /// `None` if the slice is empty. | |
53 | /// | |
54 | /// Depending on the implementation, complexity is expected to be `O(1)`. | |
55 | fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item> | |
56 | where R: Rng + ?Sized; | |
57 | ||
58 | /// Produces an iterator that chooses `amount` elements from the slice at | |
59 | /// random without repeating any, and returns them in random order. | |
60 | /// | |
61 | /// In case this API is not sufficiently flexible, use `index::sample` then | |
62 | /// apply the indices to the slice. | |
63 | /// | |
64 | /// Complexity is expected to be the same as `index::sample`. | |
65 | /// | |
66 | /// # Example | |
67 | /// ``` | |
68 | /// use rand::seq::SliceRandom; | |
69 | /// | |
70 | /// let mut rng = &mut rand::thread_rng(); | |
71 | /// let sample = "Hello, audience!".as_bytes(); | |
72 | /// | |
73 | /// // collect the results into a vector: | |
74 | /// let v: Vec<u8> = sample.choose_multiple(&mut rng, 3).cloned().collect(); | |
75 | /// | |
76 | /// // store in a buffer: | |
77 | /// let mut buf = [0u8; 5]; | |
78 | /// for (b, slot) in sample.choose_multiple(&mut rng, buf.len()).zip(buf.iter_mut()) { | |
79 | /// *slot = *b; | |
80 | /// } | |
81 | /// ``` | |
82 | #[cfg(feature = "alloc")] | |
83 | fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item> | |
84 | where R: Rng + ?Sized; | |
85 | ||
86 | /// Similar to [`choose`], where the likelihood of each outcome may be | |
87 | /// specified. The specified function `weight` maps items `x` to a relative | |
88 | /// likelihood `weight(x)`. The probability of each item being selected is | |
89 | /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`. | |
90 | /// | |
91 | /// # Example | |
92 | /// | |
93 | /// ``` | |
94 | /// use rand::prelude::*; | |
95 | /// | |
96 | /// let choices = [('a', 2), ('b', 1), ('c', 1)]; | |
97 | /// let mut rng = thread_rng(); | |
98 | /// // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c' | |
99 | /// println!("{:?}", choices.choose_weighted(&mut rng, |item| item.1).unwrap().0); | |
100 | /// ``` | |
101 | /// [`choose`]: trait.SliceRandom.html#method.choose | |
102 | #[cfg(feature = "alloc")] | |
103 | fn choose_weighted<R, F, B, X>(&self, rng: &mut R, weight: F) -> Result<&Self::Item, WeightedError> | |
104 | where R: Rng + ?Sized, | |
105 | F: Fn(&Self::Item) -> B, | |
106 | B: SampleBorrow<X>, | |
107 | X: SampleUniform + | |
108 | for<'a> ::core::ops::AddAssign<&'a X> + | |
109 | ::core::cmp::PartialOrd<X> + | |
110 | Clone + | |
111 | Default; | |
112 | ||
113 | /// Similar to [`choose_mut`], where the likelihood of each outcome may be | |
114 | /// specified. The specified function `weight` maps items `x` to a relative | |
115 | /// likelihood `weight(x)`. The probability of each item being selected is | |
116 | /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`. | |
117 | /// | |
118 | /// See also [`choose_weighted`]. | |
119 | /// | |
120 | /// [`choose_mut`]: trait.SliceRandom.html#method.choose_mut | |
121 | /// [`choose_weighted`]: trait.SliceRandom.html#method.choose_weighted | |
122 | #[cfg(feature = "alloc")] | |
123 | fn choose_weighted_mut<R, F, B, X>(&mut self, rng: &mut R, weight: F) -> Result<&mut Self::Item, WeightedError> | |
124 | where R: Rng + ?Sized, | |
125 | F: Fn(&Self::Item) -> B, | |
126 | B: SampleBorrow<X>, | |
127 | X: SampleUniform + | |
128 | for<'a> ::core::ops::AddAssign<&'a X> + | |
129 | ::core::cmp::PartialOrd<X> + | |
130 | Clone + | |
131 | Default; | |
132 | ||
133 | /// Shuffle a mutable slice in place. | |
134 | /// | |
135 | /// Depending on the implementation, complexity is expected to be `O(1)`. | |
136 | /// | |
137 | /// # Example | |
138 | /// | |
139 | /// ``` | |
140 | /// use rand::thread_rng; | |
141 | /// use rand::seq::SliceRandom; | |
142 | /// | |
143 | /// let mut rng = thread_rng(); | |
144 | /// let mut y = [1, 2, 3, 4, 5]; | |
145 | /// println!("Unshuffled: {:?}", y); | |
146 | /// y.shuffle(&mut rng); | |
147 | /// println!("Shuffled: {:?}", y); | |
148 | /// ``` | |
149 | fn shuffle<R>(&mut self, rng: &mut R) where R: Rng + ?Sized; | |
150 | ||
151 | /// Shuffle a slice in place, but exit early. | |
152 | /// | |
153 | /// Returns two mutable slices from the source slice. The first contains | |
154 | /// `amount` elements randomly permuted. The second has the remaining | |
155 | /// elements that are not fully shuffled. | |
156 | /// | |
157 | /// This is an efficient method to select `amount` elements at random from | |
158 | /// the slice, provided the slice may be mutated. | |
159 | /// | |
160 | /// If you only need to choose elements randomly and `amount > self.len()/2` | |
161 | /// then you may improve performance by taking | |
162 | /// `amount = values.len() - amount` and using only the second slice. | |
163 | /// | |
164 | /// If `amount` is greater than the number of elements in the slice, this | |
165 | /// will perform a full shuffle. | |
166 | /// | |
167 | /// Complexity is expected to be `O(m)` where `m = amount`. | |
168 | fn partial_shuffle<R>(&mut self, rng: &mut R, amount: usize) | |
169 | -> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized; | |
170 | } | |
171 | ||
172 | /// Extension trait on iterators, providing random sampling methods. | |
173 | pub trait IteratorRandom: Iterator + Sized { | |
174 | /// Choose one element at random from the iterator. If you have a slice, | |
175 | /// it's significantly faster to call the [`choose`] or [`choose_mut`] | |
176 | /// functions using the slice instead. | |
177 | /// | |
178 | /// Returns `None` if and only if the iterator is empty. | |
179 | /// | |
180 | /// Complexity is `O(n)`, where `n` is the length of the iterator. | |
181 | /// This likely consumes multiple random numbers, but the exact number | |
182 | /// is unspecified. | |
183 | /// | |
184 | /// [`choose`]: trait.SliceRandom.html#method.choose | |
185 | /// [`choose_mut`]: trait.SliceRandom.html#method.choose_mut | |
186 | fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item> | |
187 | where R: Rng + ?Sized | |
188 | { | |
189 | let (mut lower, mut upper) = self.size_hint(); | |
190 | let mut consumed = 0; | |
191 | let mut result = None; | |
192 | ||
193 | if upper == Some(lower) { | |
194 | return if lower == 0 { None } else { self.nth(rng.gen_range(0, lower)) }; | |
195 | } | |
196 | ||
197 | // Continue until the iterator is exhausted | |
198 | loop { | |
199 | if lower > 1 { | |
200 | let ix = rng.gen_range(0, lower + consumed); | |
201 | let skip; | |
202 | if ix < lower { | |
203 | result = self.nth(ix); | |
204 | skip = lower - (ix + 1); | |
205 | } else { | |
206 | skip = lower; | |
207 | } | |
208 | if upper == Some(lower) { | |
209 | return result; | |
210 | } | |
211 | consumed += lower; | |
212 | if skip > 0 { | |
213 | self.nth(skip - 1); | |
214 | } | |
215 | } else { | |
216 | let elem = self.next(); | |
217 | if elem.is_none() { | |
218 | return result; | |
219 | } | |
220 | consumed += 1; | |
221 | let denom = consumed as f64; // accurate to 2^53 elements | |
222 | if rng.gen_bool(1.0 / denom) { | |
223 | result = elem; | |
224 | } | |
225 | } | |
226 | ||
227 | let hint = self.size_hint(); | |
228 | lower = hint.0; | |
229 | upper = hint.1; | |
230 | } | |
231 | } | |
232 | ||
233 | /// Collects `amount` values at random from the iterator into a supplied | |
234 | /// buffer. | |
235 | /// | |
236 | /// Although the elements are selected randomly, the order of elements in | |
237 | /// the buffer is neither stable nor fully random. If random ordering is | |
238 | /// desired, shuffle the result. | |
239 | /// | |
240 | /// Returns the number of elements added to the buffer. This equals `amount` | |
241 | /// unless the iterator contains insufficient elements, in which case this | |
242 | /// equals the number of elements available. | |
243 | /// | |
244 | /// Complexity is `O(n)` where `n` is the length of the iterator. | |
245 | fn choose_multiple_fill<R>(mut self, rng: &mut R, buf: &mut [Self::Item]) | |
246 | -> usize where R: Rng + ?Sized | |
247 | { | |
248 | let amount = buf.len(); | |
249 | let mut len = 0; | |
250 | while len < amount { | |
251 | if let Some(elem) = self.next() { | |
252 | buf[len] = elem; | |
253 | len += 1; | |
254 | } else { | |
255 | // Iterator exhausted; stop early | |
256 | return len; | |
257 | } | |
258 | } | |
259 | ||
260 | // Continue, since the iterator was not exhausted | |
261 | for (i, elem) in self.enumerate() { | |
262 | let k = rng.gen_range(0, i + 1 + amount); | |
263 | if let Some(slot) = buf.get_mut(k) { | |
264 | *slot = elem; | |
265 | } | |
266 | } | |
267 | len | |
268 | } | |
269 | ||
270 | /// Collects `amount` values at random from the iterator into a vector. | |
271 | /// | |
272 | /// This is equivalent to `choose_multiple_fill` except for the result type. | |
273 | /// | |
274 | /// Although the elements are selected randomly, the order of elements in | |
275 | /// the buffer is neither stable nor fully random. If random ordering is | |
276 | /// desired, shuffle the result. | |
277 | /// | |
278 | /// The length of the returned vector equals `amount` unless the iterator | |
279 | /// contains insufficient elements, in which case it equals the number of | |
280 | /// elements available. | |
281 | /// | |
282 | /// Complexity is `O(n)` where `n` is the length of the iterator. | |
283 | #[cfg(feature = "alloc")] | |
284 | fn choose_multiple<R>(mut self, rng: &mut R, amount: usize) -> Vec<Self::Item> | |
285 | where R: Rng + ?Sized | |
286 | { | |
287 | let mut reservoir = Vec::with_capacity(amount); | |
288 | reservoir.extend(self.by_ref().take(amount)); | |
289 | ||
290 | // Continue unless the iterator was exhausted | |
291 | // | |
292 | // note: this prevents iterators that "restart" from causing problems. | |
293 | // If the iterator stops once, then so do we. | |
294 | if reservoir.len() == amount { | |
295 | for (i, elem) in self.enumerate() { | |
296 | let k = rng.gen_range(0, i + 1 + amount); | |
297 | if let Some(slot) = reservoir.get_mut(k) { | |
298 | *slot = elem; | |
299 | } | |
300 | } | |
301 | } else { | |
302 | // Don't hang onto extra memory. There is a corner case where | |
303 | // `amount` was much less than `self.len()`. | |
304 | reservoir.shrink_to_fit(); | |
305 | } | |
306 | reservoir | |
307 | } | |
308 | } | |
309 | ||
310 | ||
311 | impl<T> SliceRandom for [T] { | |
312 | type Item = T; | |
313 | ||
314 | fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item> | |
315 | where R: Rng + ?Sized | |
316 | { | |
317 | if self.is_empty() { | |
318 | None | |
319 | } else { | |
320 | Some(&self[rng.gen_range(0, self.len())]) | |
321 | } | |
322 | } | |
323 | ||
324 | fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item> | |
325 | where R: Rng + ?Sized | |
326 | { | |
327 | if self.is_empty() { | |
328 | None | |
329 | } else { | |
330 | let len = self.len(); | |
331 | Some(&mut self[rng.gen_range(0, len)]) | |
332 | } | |
333 | } | |
334 | ||
335 | #[cfg(feature = "alloc")] | |
336 | fn choose_multiple<R>(&self, rng: &mut R, amount: usize) | |
337 | -> SliceChooseIter<Self, Self::Item> | |
338 | where R: Rng + ?Sized | |
339 | { | |
340 | let amount = ::core::cmp::min(amount, self.len()); | |
341 | SliceChooseIter { | |
342 | slice: self, | |
343 | _phantom: Default::default(), | |
344 | indices: index::sample(rng, self.len(), amount).into_iter(), | |
345 | } | |
346 | } | |
347 | ||
348 | #[cfg(feature = "alloc")] | |
349 | fn choose_weighted<R, F, B, X>(&self, rng: &mut R, weight: F) -> Result<&Self::Item, WeightedError> | |
350 | where R: Rng + ?Sized, | |
351 | F: Fn(&Self::Item) -> B, | |
352 | B: SampleBorrow<X>, | |
353 | X: SampleUniform + | |
354 | for<'a> ::core::ops::AddAssign<&'a X> + | |
355 | ::core::cmp::PartialOrd<X> + | |
356 | Clone + | |
357 | Default { | |
358 | use distributions::{Distribution, WeightedIndex}; | |
359 | let distr = WeightedIndex::new(self.iter().map(weight))?; | |
360 | Ok(&self[distr.sample(rng)]) | |
361 | } | |
362 | ||
363 | #[cfg(feature = "alloc")] | |
364 | fn choose_weighted_mut<R, F, B, X>(&mut self, rng: &mut R, weight: F) -> Result<&mut Self::Item, WeightedError> | |
365 | where R: Rng + ?Sized, | |
366 | F: Fn(&Self::Item) -> B, | |
367 | B: SampleBorrow<X>, | |
368 | X: SampleUniform + | |
369 | for<'a> ::core::ops::AddAssign<&'a X> + | |
370 | ::core::cmp::PartialOrd<X> + | |
371 | Clone + | |
372 | Default { | |
373 | use distributions::{Distribution, WeightedIndex}; | |
374 | let distr = WeightedIndex::new(self.iter().map(weight))?; | |
375 | Ok(&mut self[distr.sample(rng)]) | |
376 | } | |
377 | ||
378 | fn shuffle<R>(&mut self, rng: &mut R) where R: Rng + ?Sized | |
379 | { | |
380 | for i in (1..self.len()).rev() { | |
381 | // invariant: elements with index > i have been locked in place. | |
382 | self.swap(i, rng.gen_range(0, i + 1)); | |
383 | } | |
384 | } | |
385 | ||
386 | fn partial_shuffle<R>(&mut self, rng: &mut R, amount: usize) | |
387 | -> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized | |
388 | { | |
389 | // This applies Durstenfeld's algorithm for the | |
390 | // [Fisher–Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm) | |
391 | // for an unbiased permutation, but exits early after choosing `amount` | |
392 | // elements. | |
393 | ||
394 | let len = self.len(); | |
395 | let end = if amount >= len { 0 } else { len - amount }; | |
396 | ||
397 | for i in (end..len).rev() { | |
398 | // invariant: elements with index > i have been locked in place. | |
399 | self.swap(i, rng.gen_range(0, i + 1)); | |
400 | } | |
401 | let r = self.split_at_mut(end); | |
402 | (r.1, r.0) | |
403 | } | |
404 | } | |
405 | ||
406 | impl<I> IteratorRandom for I where I: Iterator + Sized {} | |
407 | ||
408 | ||
409 | /// Iterator over multiple choices, as returned by [`SliceRandom::choose_multiple]( | |
410 | /// trait.SliceRandom.html#method.choose_multiple). | |
411 | #[cfg(feature = "alloc")] | |
412 | #[derive(Debug)] | |
413 | pub struct SliceChooseIter<'a, S: ?Sized + 'a, T: 'a> { | |
414 | slice: &'a S, | |
415 | _phantom: ::core::marker::PhantomData<T>, | |
416 | indices: index::IndexVecIntoIter, | |
417 | } | |
418 | ||
419 | #[cfg(feature = "alloc")] | |
420 | impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> Iterator for SliceChooseIter<'a, S, T> { | |
421 | type Item = &'a T; | |
422 | ||
423 | fn next(&mut self) -> Option<Self::Item> { | |
424 | // TODO: investigate using SliceIndex::get_unchecked when stable | |
425 | self.indices.next().map(|i| &self.slice[i as usize]) | |
426 | } | |
427 | ||
428 | fn size_hint(&self) -> (usize, Option<usize>) { | |
429 | (self.indices.len(), Some(self.indices.len())) | |
430 | } | |
431 | } | |
432 | ||
433 | #[cfg(feature = "alloc")] | |
434 | impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> ExactSizeIterator | |
435 | for SliceChooseIter<'a, S, T> | |
436 | { | |
437 | fn len(&self) -> usize { | |
438 | self.indices.len() | |
439 | } | |
440 | } | |
441 | ||
442 | ||
443 | /// Randomly sample `amount` elements from a finite iterator. | |
444 | /// | |
445 | /// Deprecated: use [`IteratorRandom::choose_multiple`] instead. | |
446 | /// | |
447 | /// [`IteratorRandom::choose_multiple`]: trait.IteratorRandom.html#method.choose_multiple | |
448 | #[cfg(feature = "alloc")] | |
449 | #[deprecated(since="0.6.0", note="use IteratorRandom::choose_multiple instead")] | |
450 | pub fn sample_iter<T, I, R>(rng: &mut R, iterable: I, amount: usize) -> Result<Vec<T>, Vec<T>> | |
451 | where I: IntoIterator<Item=T>, | |
452 | R: Rng + ?Sized, | |
453 | { | |
454 | use seq::IteratorRandom; | |
455 | let iter = iterable.into_iter(); | |
456 | let result = iter.choose_multiple(rng, amount); | |
457 | if result.len() == amount { | |
458 | Ok(result) | |
459 | } else { | |
460 | Err(result) | |
461 | } | |
462 | } | |
463 | ||
464 | /// Randomly sample exactly `amount` values from `slice`. | |
465 | /// | |
466 | /// The values are non-repeating and in random order. | |
467 | /// | |
468 | /// This implementation uses `O(amount)` time and memory. | |
469 | /// | |
470 | /// Panics if `amount > slice.len()` | |
471 | /// | |
472 | /// Deprecated: use [`SliceRandom::choose_multiple`] instead. | |
473 | /// | |
474 | /// [`SliceRandom::choose_multiple`]: trait.SliceRandom.html#method.choose_multiple | |
475 | #[cfg(feature = "alloc")] | |
476 | #[deprecated(since="0.6.0", note="use SliceRandom::choose_multiple instead")] | |
477 | pub fn sample_slice<R, T>(rng: &mut R, slice: &[T], amount: usize) -> Vec<T> | |
478 | where R: Rng + ?Sized, | |
479 | T: Clone | |
480 | { | |
481 | let indices = index::sample(rng, slice.len(), amount).into_iter(); | |
482 | ||
483 | let mut out = Vec::with_capacity(amount); | |
484 | out.extend(indices.map(|i| slice[i].clone())); | |
485 | out | |
486 | } | |
487 | ||
488 | /// Randomly sample exactly `amount` references from `slice`. | |
489 | /// | |
490 | /// The references are non-repeating and in random order. | |
491 | /// | |
492 | /// This implementation uses `O(amount)` time and memory. | |
493 | /// | |
494 | /// Panics if `amount > slice.len()` | |
495 | /// | |
496 | /// Deprecated: use [`SliceRandom::choose_multiple`] instead. | |
497 | /// | |
498 | /// [`SliceRandom::choose_multiple`]: trait.SliceRandom.html#method.choose_multiple | |
499 | #[cfg(feature = "alloc")] | |
500 | #[deprecated(since="0.6.0", note="use SliceRandom::choose_multiple instead")] | |
501 | pub fn sample_slice_ref<'a, R, T>(rng: &mut R, slice: &'a [T], amount: usize) -> Vec<&'a T> | |
502 | where R: Rng + ?Sized | |
503 | { | |
504 | let indices = index::sample(rng, slice.len(), amount).into_iter(); | |
505 | ||
506 | let mut out = Vec::with_capacity(amount); | |
507 | out.extend(indices.map(|i| &slice[i])); | |
508 | out | |
509 | } | |
510 | ||
511 | #[cfg(test)] | |
512 | mod test { | |
513 | use super::*; | |
514 | #[cfg(feature = "alloc")] use {Rng, SeedableRng}; | |
515 | #[cfg(feature = "alloc")] use rngs::SmallRng; | |
516 | #[cfg(all(feature="alloc", not(feature="std")))] | |
517 | use alloc::vec::Vec; | |
518 | ||
519 | #[test] | |
520 | fn test_slice_choose() { | |
521 | let mut r = ::test::rng(107); | |
522 | let chars = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n']; | |
523 | let mut chosen = [0i32; 14]; | |
524 | for _ in 0..1000 { | |
525 | let picked = *chars.choose(&mut r).unwrap(); | |
526 | chosen[(picked as usize) - ('a' as usize)] += 1; | |
527 | } | |
528 | for count in chosen.iter() { | |
529 | let err = *count - (1000 / (chars.len() as i32)); | |
530 | assert!(-20 <= err && err <= 20); | |
531 | } | |
532 | ||
533 | chosen.iter_mut().for_each(|x| *x = 0); | |
534 | for _ in 0..1000 { | |
535 | *chosen.choose_mut(&mut r).unwrap() += 1; | |
536 | } | |
537 | for count in chosen.iter() { | |
538 | let err = *count - (1000 / (chosen.len() as i32)); | |
539 | assert!(-20 <= err && err <= 20); | |
540 | } | |
541 | ||
542 | let mut v: [isize; 0] = []; | |
543 | assert_eq!(v.choose(&mut r), None); | |
544 | assert_eq!(v.choose_mut(&mut r), None); | |
545 | } | |
546 | ||
547 | #[derive(Clone)] | |
548 | struct UnhintedIterator<I: Iterator + Clone> { | |
549 | iter: I, | |
550 | } | |
551 | impl<I: Iterator + Clone> Iterator for UnhintedIterator<I> { | |
552 | type Item = I::Item; | |
553 | fn next(&mut self) -> Option<Self::Item> { | |
554 | self.iter.next() | |
555 | } | |
556 | } | |
557 | ||
558 | #[derive(Clone)] | |
559 | struct ChunkHintedIterator<I: ExactSizeIterator + Iterator + Clone> { | |
560 | iter: I, | |
561 | chunk_remaining: usize, | |
562 | chunk_size: usize, | |
563 | hint_total_size: bool, | |
564 | } | |
565 | impl<I: ExactSizeIterator + Iterator + Clone> Iterator for ChunkHintedIterator<I> { | |
566 | type Item = I::Item; | |
567 | fn next(&mut self) -> Option<Self::Item> { | |
568 | if self.chunk_remaining == 0 { | |
569 | self.chunk_remaining = ::core::cmp::min(self.chunk_size, | |
570 | self.iter.len()); | |
571 | } | |
572 | self.chunk_remaining = self.chunk_remaining.saturating_sub(1); | |
573 | ||
574 | self.iter.next() | |
575 | } | |
576 | fn size_hint(&self) -> (usize, Option<usize>) { | |
577 | (self.chunk_remaining, | |
578 | if self.hint_total_size { Some(self.iter.len()) } else { None }) | |
579 | } | |
580 | } | |
581 | ||
582 | #[derive(Clone)] | |
583 | struct WindowHintedIterator<I: ExactSizeIterator + Iterator + Clone> { | |
584 | iter: I, | |
585 | window_size: usize, | |
586 | hint_total_size: bool, | |
587 | } | |
588 | impl<I: ExactSizeIterator + Iterator + Clone> Iterator for WindowHintedIterator<I> { | |
589 | type Item = I::Item; | |
590 | fn next(&mut self) -> Option<Self::Item> { | |
591 | self.iter.next() | |
592 | } | |
593 | fn size_hint(&self) -> (usize, Option<usize>) { | |
594 | (::core::cmp::min(self.iter.len(), self.window_size), | |
595 | if self.hint_total_size { Some(self.iter.len()) } else { None }) | |
596 | } | |
597 | } | |
598 | ||
599 | #[test] | |
600 | fn test_iterator_choose() { | |
601 | let r = &mut ::test::rng(109); | |
602 | fn test_iter<R: Rng + ?Sized, Iter: Iterator<Item=usize> + Clone>(r: &mut R, iter: Iter) { | |
603 | let mut chosen = [0i32; 9]; | |
604 | for _ in 0..1000 { | |
605 | let picked = iter.clone().choose(r).unwrap(); | |
606 | chosen[picked] += 1; | |
607 | } | |
608 | for count in chosen.iter() { | |
609 | // Samples should follow Binomial(1000, 1/9) | |
610 | // Octave: binopdf(x, 1000, 1/9) gives the prob of *count == x | |
611 | // Note: have seen 153, which is unlikely but not impossible. | |
612 | assert!(72 < *count && *count < 154, "count not close to 1000/9: {}", count); | |
613 | } | |
614 | } | |
615 | ||
616 | test_iter(r, 0..9); | |
617 | test_iter(r, [0, 1, 2, 3, 4, 5, 6, 7, 8].iter().cloned()); | |
618 | #[cfg(feature = "alloc")] | |
619 | test_iter(r, (0..9).collect::<Vec<_>>().into_iter()); | |
620 | test_iter(r, UnhintedIterator { iter: 0..9 }); | |
621 | test_iter(r, ChunkHintedIterator { iter: 0..9, chunk_size: 4, chunk_remaining: 4, hint_total_size: false }); | |
622 | test_iter(r, ChunkHintedIterator { iter: 0..9, chunk_size: 4, chunk_remaining: 4, hint_total_size: true }); | |
623 | test_iter(r, WindowHintedIterator { iter: 0..9, window_size: 2, hint_total_size: false }); | |
624 | test_iter(r, WindowHintedIterator { iter: 0..9, window_size: 2, hint_total_size: true }); | |
625 | ||
626 | assert_eq!((0..0).choose(r), None); | |
627 | assert_eq!(UnhintedIterator{ iter: 0..0 }.choose(r), None); | |
628 | } | |
629 | ||
630 | #[test] | |
631 | fn test_shuffle() { | |
632 | let mut r = ::test::rng(108); | |
633 | let empty: &mut [isize] = &mut []; | |
634 | empty.shuffle(&mut r); | |
635 | let mut one = [1]; | |
636 | one.shuffle(&mut r); | |
637 | let b: &[_] = &[1]; | |
638 | assert_eq!(one, b); | |
639 | ||
640 | let mut two = [1, 2]; | |
641 | two.shuffle(&mut r); | |
642 | assert!(two == [1, 2] || two == [2, 1]); | |
643 | ||
644 | fn move_last(slice: &mut [usize], pos: usize) { | |
645 | // use slice[pos..].rotate_left(1); once we can use that | |
646 | let last_val = slice[pos]; | |
647 | for i in pos..slice.len() - 1 { | |
648 | slice[i] = slice[i + 1]; | |
649 | } | |
650 | *slice.last_mut().unwrap() = last_val; | |
651 | } | |
652 | let mut counts = [0i32; 24]; | |
653 | for _ in 0..10000 { | |
654 | let mut arr: [usize; 4] = [0, 1, 2, 3]; | |
655 | arr.shuffle(&mut r); | |
656 | let mut permutation = 0usize; | |
657 | let mut pos_value = counts.len(); | |
658 | for i in 0..4 { | |
659 | pos_value /= 4 - i; | |
660 | let pos = arr.iter().position(|&x| x == i).unwrap(); | |
661 | assert!(pos < (4 - i)); | |
662 | permutation += pos * pos_value; | |
663 | move_last(&mut arr, pos); | |
664 | assert_eq!(arr[3], i); | |
665 | } | |
666 | for i in 0..4 { | |
667 | assert_eq!(arr[i], i); | |
668 | } | |
669 | counts[permutation] += 1; | |
670 | } | |
671 | for count in counts.iter() { | |
672 | let err = *count - 10000i32 / 24; | |
673 | assert!(-50 <= err && err <= 50); | |
674 | } | |
675 | } | |
676 | ||
677 | #[test] | |
678 | fn test_partial_shuffle() { | |
679 | let mut r = ::test::rng(118); | |
680 | ||
681 | let mut empty: [u32; 0] = []; | |
682 | let res = empty.partial_shuffle(&mut r, 10); | |
683 | assert_eq!((res.0.len(), res.1.len()), (0, 0)); | |
684 | ||
685 | let mut v = [1, 2, 3, 4, 5]; | |
686 | let res = v.partial_shuffle(&mut r, 2); | |
687 | assert_eq!((res.0.len(), res.1.len()), (2, 3)); | |
688 | assert!(res.0[0] != res.0[1]); | |
689 | // First elements are only modified if selected, so at least one isn't modified: | |
690 | assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3); | |
691 | } | |
692 | ||
693 | #[test] | |
694 | #[cfg(feature = "alloc")] | |
695 | fn test_sample_iter() { | |
696 | let min_val = 1; | |
697 | let max_val = 100; | |
698 | ||
699 | let mut r = ::test::rng(401); | |
700 | let vals = (min_val..max_val).collect::<Vec<i32>>(); | |
701 | let small_sample = vals.iter().choose_multiple(&mut r, 5); | |
702 | let large_sample = vals.iter().choose_multiple(&mut r, vals.len() + 5); | |
703 | ||
704 | assert_eq!(small_sample.len(), 5); | |
705 | assert_eq!(large_sample.len(), vals.len()); | |
706 | // no randomization happens when amount >= len | |
707 | assert_eq!(large_sample, vals.iter().collect::<Vec<_>>()); | |
708 | ||
709 | assert!(small_sample.iter().all(|e| { | |
710 | **e >= min_val && **e <= max_val | |
711 | })); | |
712 | } | |
713 | ||
714 | #[test] | |
715 | #[cfg(feature = "alloc")] | |
716 | #[allow(deprecated)] | |
717 | fn test_sample_slice_boundaries() { | |
718 | let empty: &[u8] = &[]; | |
719 | ||
720 | let mut r = ::test::rng(402); | |
721 | ||
722 | // sample 0 items | |
723 | assert_eq!(&sample_slice(&mut r, empty, 0)[..], [0u8; 0]); | |
724 | assert_eq!(&sample_slice(&mut r, &[42, 2, 42], 0)[..], [0u8; 0]); | |
725 | ||
726 | // sample 1 item | |
727 | assert_eq!(&sample_slice(&mut r, &[42], 1)[..], [42]); | |
728 | let v = sample_slice(&mut r, &[1, 42], 1)[0]; | |
729 | assert!(v == 1 || v == 42); | |
730 | ||
731 | // sample "all" the items | |
732 | let v = sample_slice(&mut r, &[42, 133], 2); | |
733 | assert!(&v[..] == [42, 133] || v[..] == [133, 42]); | |
734 | ||
735 | // Make sure lucky 777's aren't lucky | |
736 | let slice = &[42, 777]; | |
737 | let mut num_42 = 0; | |
738 | let total = 1000; | |
739 | for _ in 0..total { | |
740 | let v = sample_slice(&mut r, slice, 1); | |
741 | assert_eq!(v.len(), 1); | |
742 | let v = v[0]; | |
743 | assert!(v == 42 || v == 777); | |
744 | if v == 42 { | |
745 | num_42 += 1; | |
746 | } | |
747 | } | |
748 | let ratio_42 = num_42 as f64 / 1000 as f64; | |
749 | assert!(0.4 <= ratio_42 || ratio_42 <= 0.6, "{}", ratio_42); | |
750 | } | |
751 | ||
752 | #[test] | |
753 | #[cfg(feature = "alloc")] | |
754 | #[allow(deprecated)] | |
755 | fn test_sample_slice() { | |
756 | let seeded_rng = SmallRng::from_seed; | |
757 | ||
758 | let mut r = ::test::rng(403); | |
759 | ||
760 | for n in 1..20 { | |
761 | let length = 5*n - 4; // 1, 6, ... | |
762 | let amount = r.gen_range(0, length); | |
763 | let mut seed = [0u8; 16]; | |
764 | r.fill(&mut seed); | |
765 | ||
766 | // assert the basics work | |
767 | let regular = index::sample(&mut seeded_rng(seed), length, amount); | |
768 | assert_eq!(regular.len(), amount); | |
769 | assert!(regular.iter().all(|e| e < length)); | |
770 | ||
771 | // also test that sampling the slice works | |
772 | let vec: Vec<u32> = (0..(length as u32)).collect(); | |
773 | let result = sample_slice(&mut seeded_rng(seed), &vec, amount); | |
774 | assert_eq!(result, regular.iter().map(|i| i as u32).collect::<Vec<_>>()); | |
775 | ||
776 | let result = sample_slice_ref(&mut seeded_rng(seed), &vec, amount); | |
777 | assert!(result.iter().zip(regular.iter()).all(|(i,j)| **i == j as u32)); | |
778 | } | |
779 | } | |
780 | ||
781 | #[test] | |
782 | #[cfg(feature = "alloc")] | |
783 | fn test_weighted() { | |
784 | let mut r = ::test::rng(406); | |
785 | const N_REPS: u32 = 3000; | |
786 | let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7]; | |
787 | let total_weight = weights.iter().sum::<u32>() as f32; | |
788 | ||
789 | let verify = |result: [i32; 14]| { | |
790 | for (i, count) in result.iter().enumerate() { | |
791 | let exp = (weights[i] * N_REPS) as f32 / total_weight; | |
792 | let mut err = (*count as f32 - exp).abs(); | |
793 | if err != 0.0 { | |
794 | err /= exp; | |
795 | } | |
796 | assert!(err <= 0.25); | |
797 | } | |
798 | }; | |
799 | ||
800 | // choose_weighted | |
801 | fn get_weight<T>(item: &(u32, T)) -> u32 { | |
802 | item.0 | |
803 | } | |
804 | let mut chosen = [0i32; 14]; | |
805 | let mut items = [(0u32, 0usize); 14]; // (weight, index) | |
806 | for (i, item) in items.iter_mut().enumerate() { | |
807 | *item = (weights[i], i); | |
808 | } | |
809 | for _ in 0..N_REPS { | |
810 | let item = items.choose_weighted(&mut r, get_weight).unwrap(); | |
811 | chosen[item.1] += 1; | |
812 | } | |
813 | verify(chosen); | |
814 | ||
815 | // choose_weighted_mut | |
816 | let mut items = [(0u32, 0i32); 14]; // (weight, count) | |
817 | for (i, item) in items.iter_mut().enumerate() { | |
818 | *item = (weights[i], 0); | |
819 | } | |
820 | for _ in 0..N_REPS { | |
821 | items.choose_weighted_mut(&mut r, get_weight).unwrap().1 += 1; | |
822 | } | |
823 | for (ch, item) in chosen.iter_mut().zip(items.iter()) { | |
824 | *ch = item.1; | |
825 | } | |
826 | verify(chosen); | |
827 | ||
828 | // Check error cases | |
829 | let empty_slice = &mut [10][0..0]; | |
830 | assert_eq!(empty_slice.choose_weighted(&mut r, |_| 1), Err(WeightedError::NoItem)); | |
831 | assert_eq!(empty_slice.choose_weighted_mut(&mut r, |_| 1), Err(WeightedError::NoItem)); | |
832 | assert_eq!(['x'].choose_weighted_mut(&mut r, |_| 0), Err(WeightedError::AllWeightsZero)); | |
833 | assert_eq!([0, -1].choose_weighted_mut(&mut r, |x| *x), Err(WeightedError::NegativeWeight)); | |
834 | assert_eq!([-1, 0].choose_weighted_mut(&mut r, |x| *x), Err(WeightedError::NegativeWeight)); | |
835 | } | |
836 | } |