Skip to content

Commit

Permalink
Add choose_multple function for choosing several values (#55)
Browse files Browse the repository at this point in the history
  • Loading branch information
notgull authored Mar 23, 2023
1 parent f939572 commit 5706126
Show file tree
Hide file tree
Showing 6 changed files with 84 additions and 3 deletions.
5 changes: 4 additions & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,9 @@ jobs:
- name: Build with no default features
# Use no-std target to ensure we don't link to std.
run: cargo build --no-default-features --target thumbv7m-none-eabi
- name: Build with no default features and alloc
# Use no-std target to ensure we don't link to std.
run: cargo build --no-default-features --features alloc --target thumbv7m-none-eabi
- name: Test wasm
run: wasm-pack test --headless --chrome
if: startsWith(matrix.os, 'ubuntu')
Expand All @@ -75,7 +78,7 @@ jobs:
matrix:
# When updating this, the reminder to update the minimum supported
# Rust version in Cargo.toml.
rust: ['1.34']
rust: ['1.36']
steps:
- uses: actions/checkout@v3
- name: Install Rust
Expand Down
5 changes: 3 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ name = "fastrand"
version = "1.9.0"
authors = ["Stjepan Glavina <stjepang@gmail.com>"]
edition = "2018"
rust-version = "1.34"
rust-version = "1.36"
description = "A simple and fast random number generator"
license = "Apache-2.0 OR MIT"
repository = "https://github.com/smol-rs/fastrand"
Expand All @@ -16,7 +16,8 @@ exclude = ["/.*"]

[features]
default = ["std"]
std = ["instant"]
alloc = []
std = ["alloc", "instant"]

[target.'cfg(all(target_arch = "wasm32", not(target_os = "wasi")))'.dependencies]
instant = { version = "0.1", optional = true }
Expand Down
7 changes: 7 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,13 @@ let i = fastrand::usize(..v.len());
let elem = v[i];
```

Sample values from an array with `O(n)` complexity (`n` is the length of array):

```rust
fastrand::choose_multiple(vec![1, 4, 5].iter(), 2);
fastrand::choose_multiple(0..20, 12);
```

Shuffle an array:

```rust
Expand Down
5 changes: 5 additions & 0 deletions src/global_rng.rs
Original file line number Diff line number Diff line change
Expand Up @@ -185,3 +185,8 @@ pub fn f32() -> f32 {
pub fn f64() -> f64 {
with_rng(|r| r.f64())
}

/// Collects `amount` values at random from the iterator into a vector.
pub fn choose_multiple<T: Iterator>(source: T, amount: usize) -> Vec<T::Item> {
with_rng(|rng| rng.choose_multiple(source, amount))
}
52 changes: 52 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,14 @@
//! let elem = v[i];
//! ```
//!
//! Sample values from an array with `O(n)` complexity (`n` is the length of array):
//!
//! ```
//! fastrand::choose_multiple(vec![1, 4, 5].iter(), 2);
//! fastrand::choose_multiple(0..20, 12);
//! ```
//!
//!
//! Shuffle an array:
//!
//! ```
Expand Down Expand Up @@ -76,9 +84,15 @@
#![forbid(unsafe_code)]
#![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]

#[cfg(feature = "alloc")]
extern crate alloc;

use core::convert::{TryFrom, TryInto};
use core::ops::{Bound, RangeBounds};

#[cfg(feature = "alloc")]
use alloc::vec::Vec;

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
mod global_rng;
Expand Down Expand Up @@ -332,6 +346,44 @@ impl Rng {
f64::from_bits((1 << (b - 2)) - (1 << f) + (self.u64(..) >> (b - f))) - 1.0
}

/// Collects `amount` values at random from the iterator into a vector.
///
/// The length of the returned vector equals `amount` unless the iterator
/// contains insufficient elements, in which case it equals the number of
/// elements available.
///
/// Complexity is `O(n)` where `n` is the length of the iterator.
#[cfg(feature = "alloc")]
#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
pub fn choose_multiple<T: Iterator>(&mut self, mut source: T, amount: usize) -> Vec<T::Item> {
// Adapted from: https://docs.rs/rand/latest/rand/seq/trait.IteratorRandom.html#method.choose_multiple
let mut reservoir = Vec::with_capacity(amount);

reservoir.extend(source.by_ref().take(amount));

// Continue unless the iterator was exhausted
//
// note: this prevents iterators that "restart" from causing problems.
// If the iterator stops once, then so do we.
if reservoir.len() == amount {
for (i, elem) in source.enumerate() {
let end = i + 1 + amount;
let k = self.usize(0..end);
if let Some(slot) = reservoir.get_mut(k) {
*slot = elem;
}
}
} else {
// If less than one third of the `Vec` was used, reallocate
// so that the unused space is not wasted. There is a corner
// case where `amount` was much less than `self.len()`.
if reservoir.capacity() > 3 * reservoir.len() {
reservoir.shrink_to_fit();
}
}
reservoir
}

rng_integer!(
i8,
u8,
Expand Down
13 changes: 13 additions & 0 deletions tests/smoke.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,19 @@ fn with_seed() {
assert_eq!(a.u64(..), b.u64(..));
}

#[test]
fn choose_multiple() {
let mut a = fastrand::Rng::new();
let mut elements = (0..20).collect::<Vec<_>>();

while !elements.is_empty() {
let chosen = a.choose_multiple(0..20, 5);
for &x in &chosen {
elements.retain(|&y| y != x);
}
}
}

#[test]
fn choice() {
let items = [1, 4, 9, 5, 2, 3, 6, 7, 8, 0];
Expand Down

0 comments on commit 5706126

Please sign in to comment.