| 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 |  | 
|---|