auto merge of #8491 : robertknight/rust/7722-reservoir_sampling, r=graydon
Fixes #7722 I had a couple of queries: - Should this return an array or an iterator? - Should this be a method on iterators or on the rng? I implemented it in RngUtils as it seemed to belong with shuffle().
This commit is contained in:
commit
435020ecc4
1 changed files with 56 additions and 0 deletions
|
@ -461,6 +461,26 @@ pub trait RngUtil {
|
||||||
* ~~~
|
* ~~~
|
||||||
*/
|
*/
|
||||||
fn shuffle_mut<T>(&mut self, values: &mut [T]);
|
fn shuffle_mut<T>(&mut self, values: &mut [T]);
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sample up to `n` values from an iterator.
|
||||||
|
*
|
||||||
|
* # Example
|
||||||
|
*
|
||||||
|
* ~~~ {.rust}
|
||||||
|
*
|
||||||
|
* use std::rand;
|
||||||
|
* use std::rand::RngUtil;
|
||||||
|
*
|
||||||
|
* fn main() {
|
||||||
|
* let mut rng = rand::rng();
|
||||||
|
* let vals = range(1, 100).to_owned_vec();
|
||||||
|
* let sample = rng.sample(vals.iter(), 5);
|
||||||
|
* printfln!(sample);
|
||||||
|
* }
|
||||||
|
* ~~~
|
||||||
|
*/
|
||||||
|
fn sample<A, T: Iterator<A>>(&mut self, iter: T, n: uint) -> ~[A];
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Extension methods for random number generators
|
/// Extension methods for random number generators
|
||||||
|
@ -607,6 +627,23 @@ impl<R: Rng> RngUtil for R {
|
||||||
values.swap(i, self.gen_uint_range(0u, i + 1u));
|
values.swap(i, self.gen_uint_range(0u, i + 1u));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Randomly sample up to `n` elements from an iterator
|
||||||
|
fn sample<A, T: Iterator<A>>(&mut self, iter: T, n: uint) -> ~[A] {
|
||||||
|
let mut reservoir : ~[A] = vec::with_capacity(n);
|
||||||
|
for (i, elem) in iter.enumerate() {
|
||||||
|
if i < n {
|
||||||
|
reservoir.push(elem);
|
||||||
|
loop
|
||||||
|
}
|
||||||
|
|
||||||
|
let k = self.gen_uint_range(0, i + 1);
|
||||||
|
if k < reservoir.len() {
|
||||||
|
reservoir[k] = elem
|
||||||
|
}
|
||||||
|
}
|
||||||
|
reservoir
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create a random number generator with a default algorithm and seed.
|
/// Create a random number generator with a default algorithm and seed.
|
||||||
|
@ -914,6 +951,7 @@ pub fn random<T: Rand>() -> T {
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod test {
|
mod test {
|
||||||
|
use iterator::{Iterator, range};
|
||||||
use option::{Option, Some};
|
use option::{Option, Some};
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
@ -1130,6 +1168,24 @@ mod test {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sample() {
|
||||||
|
let MIN_VAL = 1;
|
||||||
|
let MAX_VAL = 100;
|
||||||
|
|
||||||
|
let mut r = rng();
|
||||||
|
let vals = range(MIN_VAL, MAX_VAL).to_owned_vec();
|
||||||
|
let small_sample = r.sample(vals.iter(), 5);
|
||||||
|
let large_sample = r.sample(vals.iter(), vals.len() + 5);
|
||||||
|
|
||||||
|
assert_eq!(small_sample.len(), 5);
|
||||||
|
assert_eq!(large_sample.len(), vals.len());
|
||||||
|
|
||||||
|
assert!(small_sample.iter().all(|e| {
|
||||||
|
**e >= MIN_VAL && **e <= MAX_VAL
|
||||||
|
}));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue