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 | //! Low-level API for sampling indices |
10 | |
11 | #[cfg (feature = "alloc" )] use core::slice; |
12 | |
13 | #[cfg (feature = "alloc" )] use alloc::vec::{self, Vec}; |
14 | // BTreeMap is not as fast in tests, but better than nothing. |
15 | #[cfg (all(feature = "alloc" , not(feature = "std" )))] |
16 | use alloc::collections::BTreeSet; |
17 | #[cfg (feature = "std" )] use std::collections::HashSet; |
18 | |
19 | #[cfg (feature = "std" )] |
20 | use crate::distributions::WeightedError; |
21 | |
22 | #[cfg (feature = "alloc" )] |
23 | use crate::{Rng, distributions::{uniform::SampleUniform, Distribution, Uniform}}; |
24 | |
25 | #[cfg (feature = "serde1" )] |
26 | use serde::{Serialize, Deserialize}; |
27 | |
28 | /// A vector of indices. |
29 | /// |
30 | /// Multiple internal representations are possible. |
31 | #[derive (Clone, Debug)] |
32 | #[cfg_attr (feature = "serde1" , derive(Serialize, Deserialize))] |
33 | pub enum IndexVec { |
34 | #[doc (hidden)] |
35 | U32(Vec<u32>), |
36 | #[doc (hidden)] |
37 | USize(Vec<usize>), |
38 | } |
39 | |
40 | impl IndexVec { |
41 | /// Returns the number of indices |
42 | #[inline ] |
43 | pub fn len(&self) -> usize { |
44 | match *self { |
45 | IndexVec::U32(ref v) => v.len(), |
46 | IndexVec::USize(ref v) => v.len(), |
47 | } |
48 | } |
49 | |
50 | /// Returns `true` if the length is 0. |
51 | #[inline ] |
52 | pub fn is_empty(&self) -> bool { |
53 | match *self { |
54 | IndexVec::U32(ref v) => v.is_empty(), |
55 | IndexVec::USize(ref v) => v.is_empty(), |
56 | } |
57 | } |
58 | |
59 | /// Return the value at the given `index`. |
60 | /// |
61 | /// (Note: we cannot implement [`std::ops::Index`] because of lifetime |
62 | /// restrictions.) |
63 | #[inline ] |
64 | pub fn index(&self, index: usize) -> usize { |
65 | match *self { |
66 | IndexVec::U32(ref v) => v[index] as usize, |
67 | IndexVec::USize(ref v) => v[index], |
68 | } |
69 | } |
70 | |
71 | /// Return result as a `Vec<usize>`. Conversion may or may not be trivial. |
72 | #[inline ] |
73 | pub fn into_vec(self) -> Vec<usize> { |
74 | match self { |
75 | IndexVec::U32(v) => v.into_iter().map(|i| i as usize).collect(), |
76 | IndexVec::USize(v) => v, |
77 | } |
78 | } |
79 | |
80 | /// Iterate over the indices as a sequence of `usize` values |
81 | #[inline ] |
82 | pub fn iter(&self) -> IndexVecIter<'_> { |
83 | match *self { |
84 | IndexVec::U32(ref v) => IndexVecIter::U32(v.iter()), |
85 | IndexVec::USize(ref v) => IndexVecIter::USize(v.iter()), |
86 | } |
87 | } |
88 | } |
89 | |
90 | impl IntoIterator for IndexVec { |
91 | type Item = usize; |
92 | type IntoIter = IndexVecIntoIter; |
93 | |
94 | /// Convert into an iterator over the indices as a sequence of `usize` values |
95 | #[inline ] |
96 | fn into_iter(self) -> IndexVecIntoIter { |
97 | match self { |
98 | IndexVec::U32(v: Vec) => IndexVecIntoIter::U32(v.into_iter()), |
99 | IndexVec::USize(v: Vec) => IndexVecIntoIter::USize(v.into_iter()), |
100 | } |
101 | } |
102 | } |
103 | |
104 | impl PartialEq for IndexVec { |
105 | fn eq(&self, other: &IndexVec) -> bool { |
106 | use self::IndexVec::*; |
107 | match (self, other) { |
108 | (&U32(ref v1: &Vec), &U32(ref v2: &Vec)) => v1 == v2, |
109 | (&USize(ref v1: &Vec), &USize(ref v2: &Vec)) => v1 == v2, |
110 | (&U32(ref v1: &Vec), &USize(ref v2: &Vec)) => { |
111 | (v1.len() == v2.len()) && (v1.iter().zip(v2.iter()).all(|(x: &u32, y: &usize)| *x as usize == *y)) |
112 | } |
113 | (&USize(ref v1: &Vec), &U32(ref v2: &Vec)) => { |
114 | (v1.len() == v2.len()) && (v1.iter().zip(v2.iter()).all(|(x: &usize, y: &u32)| *x == *y as usize)) |
115 | } |
116 | } |
117 | } |
118 | } |
119 | |
120 | impl From<Vec<u32>> for IndexVec { |
121 | #[inline ] |
122 | fn from(v: Vec<u32>) -> Self { |
123 | IndexVec::U32(v) |
124 | } |
125 | } |
126 | |
127 | impl From<Vec<usize>> for IndexVec { |
128 | #[inline ] |
129 | fn from(v: Vec<usize>) -> Self { |
130 | IndexVec::USize(v) |
131 | } |
132 | } |
133 | |
134 | /// Return type of `IndexVec::iter`. |
135 | #[derive (Debug)] |
136 | pub enum IndexVecIter<'a> { |
137 | #[doc (hidden)] |
138 | U32(slice::Iter<'a, u32>), |
139 | #[doc (hidden)] |
140 | USize(slice::Iter<'a, usize>), |
141 | } |
142 | |
143 | impl<'a> Iterator for IndexVecIter<'a> { |
144 | type Item = usize; |
145 | |
146 | #[inline ] |
147 | fn next(&mut self) -> Option<usize> { |
148 | use self::IndexVecIter::*; |
149 | match *self { |
150 | U32(ref mut iter: &mut Iter<'_, u32>) => iter.next().map(|i: &u32| *i as usize), |
151 | USize(ref mut iter: &mut Iter<'_, usize>) => iter.next().cloned(), |
152 | } |
153 | } |
154 | |
155 | #[inline ] |
156 | fn size_hint(&self) -> (usize, Option<usize>) { |
157 | match *self { |
158 | IndexVecIter::U32(ref v: &Iter<'_, u32>) => v.size_hint(), |
159 | IndexVecIter::USize(ref v: &Iter<'_, usize>) => v.size_hint(), |
160 | } |
161 | } |
162 | } |
163 | |
164 | impl<'a> ExactSizeIterator for IndexVecIter<'a> {} |
165 | |
166 | /// Return type of `IndexVec::into_iter`. |
167 | #[derive (Clone, Debug)] |
168 | pub enum IndexVecIntoIter { |
169 | #[doc (hidden)] |
170 | U32(vec::IntoIter<u32>), |
171 | #[doc (hidden)] |
172 | USize(vec::IntoIter<usize>), |
173 | } |
174 | |
175 | impl Iterator for IndexVecIntoIter { |
176 | type Item = usize; |
177 | |
178 | #[inline ] |
179 | fn next(&mut self) -> Option<Self::Item> { |
180 | use self::IndexVecIntoIter::*; |
181 | match *self { |
182 | U32(ref mut v: &mut IntoIter) => v.next().map(|i: u32| i as usize), |
183 | USize(ref mut v: &mut IntoIter) => v.next(), |
184 | } |
185 | } |
186 | |
187 | #[inline ] |
188 | fn size_hint(&self) -> (usize, Option<usize>) { |
189 | use self::IndexVecIntoIter::*; |
190 | match *self { |
191 | U32(ref v: &IntoIter) => v.size_hint(), |
192 | USize(ref v: &IntoIter) => v.size_hint(), |
193 | } |
194 | } |
195 | } |
196 | |
197 | impl ExactSizeIterator for IndexVecIntoIter {} |
198 | |
199 | |
200 | /// Randomly sample exactly `amount` distinct indices from `0..length`, and |
201 | /// return them in random order (fully shuffled). |
202 | /// |
203 | /// This method is used internally by the slice sampling methods, but it can |
204 | /// sometimes be useful to have the indices themselves so this is provided as |
205 | /// an alternative. |
206 | /// |
207 | /// The implementation used is not specified; we automatically select the |
208 | /// fastest available algorithm for the `length` and `amount` parameters |
209 | /// (based on detailed profiling on an Intel Haswell CPU). Roughly speaking, |
210 | /// complexity is `O(amount)`, except that when `amount` is small, performance |
211 | /// is closer to `O(amount^2)`, and when `length` is close to `amount` then |
212 | /// `O(length)`. |
213 | /// |
214 | /// Note that performance is significantly better over `u32` indices than over |
215 | /// `u64` indices. Because of this we hide the underlying type behind an |
216 | /// abstraction, `IndexVec`. |
217 | /// |
218 | /// If an allocation-free `no_std` function is required, it is suggested |
219 | /// to adapt the internal `sample_floyd` implementation. |
220 | /// |
221 | /// Panics if `amount > length`. |
222 | pub fn sample<R>(rng: &mut R, length: usize, amount: usize) -> IndexVec |
223 | where R: Rng + ?Sized { |
224 | if amount > length { |
225 | panic!("`amount` of samples must be less than or equal to `length`" ); |
226 | } |
227 | if length > (::core::u32::MAX as usize) { |
228 | // We never want to use inplace here, but could use floyd's alg |
229 | // Lazy version: always use the cache alg. |
230 | return sample_rejection(rng, length, amount); |
231 | } |
232 | let amount = amount as u32; |
233 | let length = length as u32; |
234 | |
235 | // Choice of algorithm here depends on both length and amount. See: |
236 | // https://github.com/rust-random/rand/pull/479 |
237 | // We do some calculations with f32. Accuracy is not very important. |
238 | |
239 | if amount < 163 { |
240 | const C: [[f32; 2]; 2] = [[1.6, 8.0 / 45.0], [10.0, 70.0 / 9.0]]; |
241 | let j = if length < 500_000 { 0 } else { 1 }; |
242 | let amount_fp = amount as f32; |
243 | let m4 = C[0][j] * amount_fp; |
244 | // Short-cut: when amount < 12, floyd's is always faster |
245 | if amount > 11 && (length as f32) < (C[1][j] + m4) * amount_fp { |
246 | sample_inplace(rng, length, amount) |
247 | } else { |
248 | sample_floyd(rng, length, amount) |
249 | } |
250 | } else { |
251 | const C: [f32; 2] = [270.0, 330.0 / 9.0]; |
252 | let j = if length < 500_000 { 0 } else { 1 }; |
253 | if (length as f32) < C[j] * (amount as f32) { |
254 | sample_inplace(rng, length, amount) |
255 | } else { |
256 | sample_rejection(rng, length, amount) |
257 | } |
258 | } |
259 | } |
260 | |
261 | /// Randomly sample exactly `amount` distinct indices from `0..length`, and |
262 | /// return them in an arbitrary order (there is no guarantee of shuffling or |
263 | /// ordering). The weights are to be provided by the input function `weights`, |
264 | /// which will be called once for each index. |
265 | /// |
266 | /// This method is used internally by the slice sampling methods, but it can |
267 | /// sometimes be useful to have the indices themselves so this is provided as |
268 | /// an alternative. |
269 | /// |
270 | /// This implementation uses `O(length + amount)` space and `O(length)` time |
271 | /// if the "nightly" feature is enabled, or `O(length)` space and |
272 | /// `O(length + amount * log length)` time otherwise. |
273 | /// |
274 | /// Panics if `amount > length`. |
275 | #[cfg (feature = "std" )] |
276 | #[cfg_attr (doc_cfg, doc(cfg(feature = "std" )))] |
277 | pub fn sample_weighted<R, F, X>( |
278 | rng: &mut R, length: usize, weight: F, amount: usize, |
279 | ) -> Result<IndexVec, WeightedError> |
280 | where |
281 | R: Rng + ?Sized, |
282 | F: Fn(usize) -> X, |
283 | X: Into<f64>, |
284 | { |
285 | if length > (core::u32::MAX as usize) { |
286 | sample_efraimidis_spirakis(rng, length, weight, amount) |
287 | } else { |
288 | assert!(amount <= core::u32::MAX as usize); |
289 | let amount: u32 = amount as u32; |
290 | let length: u32 = length as u32; |
291 | sample_efraimidis_spirakis(rng, length, weight, amount) |
292 | } |
293 | } |
294 | |
295 | |
296 | /// Randomly sample exactly `amount` distinct indices from `0..length`, and |
297 | /// return them in an arbitrary order (there is no guarantee of shuffling or |
298 | /// ordering). The weights are to be provided by the input function `weights`, |
299 | /// which will be called once for each index. |
300 | /// |
301 | /// This implementation uses the algorithm described by Efraimidis and Spirakis |
302 | /// in this paper: https://doi.org/10.1016/j.ipl.2005.11.003 |
303 | /// It uses `O(length + amount)` space and `O(length)` time if the |
304 | /// "nightly" feature is enabled, or `O(length)` space and `O(length |
305 | /// + amount * log length)` time otherwise. |
306 | /// |
307 | /// Panics if `amount > length`. |
308 | #[cfg (feature = "std" )] |
309 | fn sample_efraimidis_spirakis<R, F, X, N>( |
310 | rng: &mut R, length: N, weight: F, amount: N, |
311 | ) -> Result<IndexVec, WeightedError> |
312 | where |
313 | R: Rng + ?Sized, |
314 | F: Fn(usize) -> X, |
315 | X: Into<f64>, |
316 | N: UInt, |
317 | IndexVec: From<Vec<N>>, |
318 | { |
319 | if amount == N::zero() { |
320 | return Ok(IndexVec::U32(Vec::new())); |
321 | } |
322 | |
323 | if amount > length { |
324 | panic!("`amount` of samples must be less than or equal to `length`" ); |
325 | } |
326 | |
327 | struct Element<N> { |
328 | index: N, |
329 | key: f64, |
330 | } |
331 | impl<N> PartialOrd for Element<N> { |
332 | fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> { |
333 | self.key.partial_cmp(&other.key) |
334 | } |
335 | } |
336 | impl<N> Ord for Element<N> { |
337 | fn cmp(&self, other: &Self) -> core::cmp::Ordering { |
338 | // partial_cmp will always produce a value, |
339 | // because we check that the weights are not nan |
340 | self.partial_cmp(other).unwrap() |
341 | } |
342 | } |
343 | impl<N> PartialEq for Element<N> { |
344 | fn eq(&self, other: &Self) -> bool { |
345 | self.key == other.key |
346 | } |
347 | } |
348 | impl<N> Eq for Element<N> {} |
349 | |
350 | #[cfg (feature = "nightly" )] |
351 | { |
352 | let mut candidates = Vec::with_capacity(length.as_usize()); |
353 | let mut index = N::zero(); |
354 | while index < length { |
355 | let weight = weight(index.as_usize()).into(); |
356 | if !(weight >= 0.) { |
357 | return Err(WeightedError::InvalidWeight); |
358 | } |
359 | |
360 | let key = rng.gen::<f64>().powf(1.0 / weight); |
361 | candidates.push(Element { index, key }); |
362 | |
363 | index += N::one(); |
364 | } |
365 | |
366 | // Partially sort the array to find the `amount` elements with the greatest |
367 | // keys. Do this by using `select_nth_unstable` to put the elements with |
368 | // the *smallest* keys at the beginning of the list in `O(n)` time, which |
369 | // provides equivalent information about the elements with the *greatest* keys. |
370 | let (_, mid, greater) |
371 | = candidates.select_nth_unstable(length.as_usize() - amount.as_usize()); |
372 | |
373 | let mut result: Vec<N> = Vec::with_capacity(amount.as_usize()); |
374 | result.push(mid.index); |
375 | for element in greater { |
376 | result.push(element.index); |
377 | } |
378 | Ok(IndexVec::from(result)) |
379 | } |
380 | |
381 | #[cfg (not(feature = "nightly" ))] |
382 | { |
383 | use alloc::collections::BinaryHeap; |
384 | |
385 | // Partially sort the array such that the `amount` elements with the largest |
386 | // keys are first using a binary max heap. |
387 | let mut candidates = BinaryHeap::with_capacity(length.as_usize()); |
388 | let mut index = N::zero(); |
389 | while index < length { |
390 | let weight = weight(index.as_usize()).into(); |
391 | if !(weight >= 0.) { |
392 | return Err(WeightedError::InvalidWeight); |
393 | } |
394 | |
395 | let key = rng.gen::<f64>().powf(1.0 / weight); |
396 | candidates.push(Element { index, key }); |
397 | |
398 | index += N::one(); |
399 | } |
400 | |
401 | let mut result: Vec<N> = Vec::with_capacity(amount.as_usize()); |
402 | while result.len() < amount.as_usize() { |
403 | result.push(candidates.pop().unwrap().index); |
404 | } |
405 | Ok(IndexVec::from(result)) |
406 | } |
407 | } |
408 | |
409 | /// Randomly sample exactly `amount` indices from `0..length`, using Floyd's |
410 | /// combination algorithm. |
411 | /// |
412 | /// The output values are fully shuffled. (Overhead is under 50%.) |
413 | /// |
414 | /// This implementation uses `O(amount)` memory and `O(amount^2)` time. |
415 | fn sample_floyd<R>(rng: &mut R, length: u32, amount: u32) -> IndexVec |
416 | where R: Rng + ?Sized { |
417 | // For small amount we use Floyd's fully-shuffled variant. For larger |
418 | // amounts this is slow due to Vec::insert performance, so we shuffle |
419 | // afterwards. Benchmarks show little overhead from extra logic. |
420 | let floyd_shuffle = amount < 50; |
421 | |
422 | debug_assert!(amount <= length); |
423 | let mut indices = Vec::with_capacity(amount as usize); |
424 | for j in length - amount..length { |
425 | let t = rng.gen_range(0..=j); |
426 | if floyd_shuffle { |
427 | if let Some(pos) = indices.iter().position(|&x| x == t) { |
428 | indices.insert(pos, j); |
429 | continue; |
430 | } |
431 | } else if indices.contains(&t) { |
432 | indices.push(j); |
433 | continue; |
434 | } |
435 | indices.push(t); |
436 | } |
437 | if !floyd_shuffle { |
438 | // Reimplement SliceRandom::shuffle with smaller indices |
439 | for i in (1..amount).rev() { |
440 | // invariant: elements with index > i have been locked in place. |
441 | indices.swap(i as usize, rng.gen_range(0..=i) as usize); |
442 | } |
443 | } |
444 | IndexVec::from(indices) |
445 | } |
446 | |
447 | /// Randomly sample exactly `amount` indices from `0..length`, using an inplace |
448 | /// partial Fisher-Yates method. |
449 | /// Sample an amount of indices using an inplace partial fisher yates method. |
450 | /// |
451 | /// This allocates the entire `length` of indices and randomizes only the first `amount`. |
452 | /// It then truncates to `amount` and returns. |
453 | /// |
454 | /// This method is not appropriate for large `length` and potentially uses a lot |
455 | /// of memory; because of this we only implement for `u32` index (which improves |
456 | /// performance in all cases). |
457 | /// |
458 | /// Set-up is `O(length)` time and memory and shuffling is `O(amount)` time. |
459 | fn sample_inplace<R>(rng: &mut R, length: u32, amount: u32) -> IndexVec |
460 | where R: Rng + ?Sized { |
461 | debug_assert!(amount <= length); |
462 | let mut indices: Vec<u32> = Vec::with_capacity(length as usize); |
463 | indices.extend(iter:0..length); |
464 | for i: u32 in 0..amount { |
465 | let j: u32 = rng.gen_range(i..length); |
466 | indices.swap(a:i as usize, b:j as usize); |
467 | } |
468 | indices.truncate(len:amount as usize); |
469 | debug_assert_eq!(indices.len(), amount as usize); |
470 | IndexVec::from(indices) |
471 | } |
472 | |
473 | trait UInt: Copy + PartialOrd + Ord + PartialEq + Eq + SampleUniform |
474 | + core::hash::Hash + core::ops::AddAssign { |
475 | fn zero() -> Self; |
476 | fn one() -> Self; |
477 | fn as_usize(self) -> usize; |
478 | } |
479 | impl UInt for u32 { |
480 | #[inline ] |
481 | fn zero() -> Self { |
482 | 0 |
483 | } |
484 | |
485 | #[inline ] |
486 | fn one() -> Self { |
487 | 1 |
488 | } |
489 | |
490 | #[inline ] |
491 | fn as_usize(self) -> usize { |
492 | self as usize |
493 | } |
494 | } |
495 | impl UInt for usize { |
496 | #[inline ] |
497 | fn zero() -> Self { |
498 | 0 |
499 | } |
500 | |
501 | #[inline ] |
502 | fn one() -> Self { |
503 | 1 |
504 | } |
505 | |
506 | #[inline ] |
507 | fn as_usize(self) -> usize { |
508 | self |
509 | } |
510 | } |
511 | |
512 | /// Randomly sample exactly `amount` indices from `0..length`, using rejection |
513 | /// sampling. |
514 | /// |
515 | /// Since `amount <<< length` there is a low chance of a random sample in |
516 | /// `0..length` being a duplicate. We test for duplicates and resample where |
517 | /// necessary. The algorithm is `O(amount)` time and memory. |
518 | /// |
519 | /// This function is generic over X primarily so that results are value-stable |
520 | /// over 32-bit and 64-bit platforms. |
521 | fn sample_rejection<X: UInt, R>(rng: &mut R, length: X, amount: X) -> IndexVec |
522 | where |
523 | R: Rng + ?Sized, |
524 | IndexVec: From<Vec<X>>, |
525 | { |
526 | debug_assert!(amount < length); |
527 | #[cfg (feature = "std" )] |
528 | let mut cache: HashSet = HashSet::with_capacity(amount.as_usize()); |
529 | #[cfg (not(feature = "std" ))] |
530 | let mut cache = BTreeSet::new(); |
531 | let distr: Uniform = Uniform::new(X::zero(), high:length); |
532 | let mut indices: Vec = Vec::with_capacity(amount.as_usize()); |
533 | for _ in 0..amount.as_usize() { |
534 | let mut pos: X = distr.sample(rng); |
535 | while !cache.insert(pos) { |
536 | pos = distr.sample(rng); |
537 | } |
538 | indices.push(pos); |
539 | } |
540 | |
541 | debug_assert_eq!(indices.len(), amount.as_usize()); |
542 | IndexVec::from(indices) |
543 | } |
544 | |
545 | #[cfg (test)] |
546 | mod test { |
547 | use super::*; |
548 | |
549 | #[test ] |
550 | #[cfg (feature = "serde1" )] |
551 | fn test_serialization_index_vec() { |
552 | let some_index_vec = IndexVec::from(vec![254_usize, 234, 2, 1]); |
553 | let de_some_index_vec: IndexVec = bincode::deserialize(&bincode::serialize(&some_index_vec).unwrap()).unwrap(); |
554 | match (some_index_vec, de_some_index_vec) { |
555 | (IndexVec::U32(a), IndexVec::U32(b)) => { |
556 | assert_eq!(a, b); |
557 | }, |
558 | (IndexVec::USize(a), IndexVec::USize(b)) => { |
559 | assert_eq!(a, b); |
560 | }, |
561 | _ => {panic!("failed to seralize/deserialize `IndexVec`" )} |
562 | } |
563 | } |
564 | |
565 | #[cfg (feature = "alloc" )] use alloc::vec; |
566 | |
567 | #[test ] |
568 | fn test_sample_boundaries() { |
569 | let mut r = crate::test::rng(404); |
570 | |
571 | assert_eq!(sample_inplace(&mut r, 0, 0).len(), 0); |
572 | assert_eq!(sample_inplace(&mut r, 1, 0).len(), 0); |
573 | assert_eq!(sample_inplace(&mut r, 1, 1).into_vec(), vec![0]); |
574 | |
575 | assert_eq!(sample_rejection(&mut r, 1u32, 0).len(), 0); |
576 | |
577 | assert_eq!(sample_floyd(&mut r, 0, 0).len(), 0); |
578 | assert_eq!(sample_floyd(&mut r, 1, 0).len(), 0); |
579 | assert_eq!(sample_floyd(&mut r, 1, 1).into_vec(), vec![0]); |
580 | |
581 | // These algorithms should be fast with big numbers. Test average. |
582 | let sum: usize = sample_rejection(&mut r, 1 << 25, 10u32).into_iter().sum(); |
583 | assert!(1 << 25 < sum && sum < (1 << 25) * 25); |
584 | |
585 | let sum: usize = sample_floyd(&mut r, 1 << 25, 10).into_iter().sum(); |
586 | assert!(1 << 25 < sum && sum < (1 << 25) * 25); |
587 | } |
588 | |
589 | #[test ] |
590 | #[cfg_attr (miri, ignore)] // Miri is too slow |
591 | fn test_sample_alg() { |
592 | let seed_rng = crate::test::rng; |
593 | |
594 | // We can't test which algorithm is used directly, but Floyd's alg |
595 | // should produce different results from the others. (Also, `inplace` |
596 | // and `cached` currently use different sizes thus produce different results.) |
597 | |
598 | // A small length and relatively large amount should use inplace |
599 | let (length, amount): (usize, usize) = (100, 50); |
600 | let v1 = sample(&mut seed_rng(420), length, amount); |
601 | let v2 = sample_inplace(&mut seed_rng(420), length as u32, amount as u32); |
602 | assert!(v1.iter().all(|e| e < length)); |
603 | assert_eq!(v1, v2); |
604 | |
605 | // Test Floyd's alg does produce different results |
606 | let v3 = sample_floyd(&mut seed_rng(420), length as u32, amount as u32); |
607 | assert!(v1 != v3); |
608 | |
609 | // A large length and small amount should use Floyd |
610 | let (length, amount): (usize, usize) = (1 << 20, 50); |
611 | let v1 = sample(&mut seed_rng(421), length, amount); |
612 | let v2 = sample_floyd(&mut seed_rng(421), length as u32, amount as u32); |
613 | assert!(v1.iter().all(|e| e < length)); |
614 | assert_eq!(v1, v2); |
615 | |
616 | // A large length and larger amount should use cache |
617 | let (length, amount): (usize, usize) = (1 << 20, 600); |
618 | let v1 = sample(&mut seed_rng(422), length, amount); |
619 | let v2 = sample_rejection(&mut seed_rng(422), length as u32, amount as u32); |
620 | assert!(v1.iter().all(|e| e < length)); |
621 | assert_eq!(v1, v2); |
622 | } |
623 | |
624 | #[cfg (feature = "std" )] |
625 | #[test ] |
626 | fn test_sample_weighted() { |
627 | let seed_rng = crate::test::rng; |
628 | for &(amount, len) in &[(0, 10), (5, 10), (10, 10)] { |
629 | let v = sample_weighted(&mut seed_rng(423), len, |i| i as f64, amount).unwrap(); |
630 | match v { |
631 | IndexVec::U32(mut indices) => { |
632 | assert_eq!(indices.len(), amount); |
633 | indices.sort_unstable(); |
634 | indices.dedup(); |
635 | assert_eq!(indices.len(), amount); |
636 | for &i in &indices { |
637 | assert!((i as usize) < len); |
638 | } |
639 | }, |
640 | IndexVec::USize(_) => panic!("expected `IndexVec::U32`" ), |
641 | } |
642 | } |
643 | } |
644 | |
645 | #[test ] |
646 | fn value_stability_sample() { |
647 | let do_test = |length, amount, values: &[u32]| { |
648 | let mut buf = [0u32; 8]; |
649 | let mut rng = crate::test::rng(410); |
650 | |
651 | let res = sample(&mut rng, length, amount); |
652 | let len = res.len().min(buf.len()); |
653 | for (x, y) in res.into_iter().zip(buf.iter_mut()) { |
654 | *y = x as u32; |
655 | } |
656 | assert_eq!( |
657 | &buf[0..len], |
658 | values, |
659 | "failed sampling {}, {}" , |
660 | length, |
661 | amount |
662 | ); |
663 | }; |
664 | |
665 | do_test(10, 6, &[8, 0, 3, 5, 9, 6]); // floyd |
666 | do_test(25, 10, &[18, 15, 14, 9, 0, 13, 5, 24]); // floyd |
667 | do_test(300, 8, &[30, 283, 150, 1, 73, 13, 285, 35]); // floyd |
668 | do_test(300, 80, &[31, 289, 248, 154, 5, 78, 19, 286]); // inplace |
669 | do_test(300, 180, &[31, 289, 248, 154, 5, 78, 19, 286]); // inplace |
670 | |
671 | do_test(1_000_000, 8, &[ |
672 | 103717, 963485, 826422, 509101, 736394, 807035, 5327, 632573, |
673 | ]); // floyd |
674 | do_test(1_000_000, 180, &[ |
675 | 103718, 963490, 826426, 509103, 736396, 807036, 5327, 632573, |
676 | ]); // rejection |
677 | } |
678 | } |
679 | |