98 lines
2.4 KiB
Rust
98 lines
2.4 KiB
Rust
use std;
|
|
import std::rand;
|
|
import vec;
|
|
|
|
// random uint less than n
|
|
fn under(r : rand::rng, n : uint) -> uint {
|
|
assert n != 0u; r.next() as uint % n
|
|
}
|
|
|
|
// random choice from a vec
|
|
fn choice<T: copy>(r : rand::rng, v : [T]) -> T {
|
|
assert vec::len(v) != 0u; v[under(r, vec::len(v))]
|
|
}
|
|
|
|
// 1 in n chance of being true
|
|
fn unlikely(r : rand::rng, n : uint) -> bool { under(r, n) == 0u }
|
|
|
|
// shuffle a vec in place
|
|
fn shuffle<T>(r : rand::rng, &v : [mut T]) {
|
|
let i = vec::len(v);
|
|
while i >= 2u {
|
|
// Loop invariant: elements with index >= i have been locked in place.
|
|
i -= 1u;
|
|
vec::swap(v, i, under(r, i + 1u)); // Lock element i in place.
|
|
}
|
|
}
|
|
|
|
// create a shuffled copy of a vec
|
|
fn shuffled<T: copy>(r : rand::rng, v : [T]) -> [T] {
|
|
let w = vec::to_mut(v);
|
|
shuffle(r, w);
|
|
vec::from_mut(w) // Shouldn't this happen automatically?
|
|
}
|
|
|
|
// sample from a population without replacement
|
|
//fn sample<T>(r : rand::rng, pop : [T], k : uint) -> [T] { fail }
|
|
|
|
// Two ways to make a weighted choice.
|
|
// * weighted_choice is O(number of choices) time
|
|
// * weighted_vec is O(total weight) space
|
|
type weighted<T> = { weight: uint, item: T };
|
|
fn weighted_choice<T: copy>(r : rand::rng, v : [weighted<T>]) -> T {
|
|
assert vec::len(v) != 0u;
|
|
let total = 0u;
|
|
for {weight: weight, item: _} in v {
|
|
total += weight;
|
|
}
|
|
assert total >= 0u;
|
|
let chosen = under(r, total);
|
|
let so_far = 0u;
|
|
for {weight: weight, item: item} in v {
|
|
so_far += weight;
|
|
if so_far > chosen {
|
|
ret item;
|
|
}
|
|
}
|
|
core::unreachable();
|
|
}
|
|
|
|
fn weighted_vec<T: copy>(v : [weighted<T>]) -> [T] {
|
|
let r = [];
|
|
for {weight: weight, item: item} in v {
|
|
let i = 0u;
|
|
while i < weight {
|
|
r += [item];
|
|
i += 1u;
|
|
}
|
|
}
|
|
r
|
|
}
|
|
|
|
fn main()
|
|
{
|
|
let r = rand::mk_rng();
|
|
|
|
log(error, under(r, 5u));
|
|
log(error, choice(r, [10, 20, 30]));
|
|
log(error, if unlikely(r, 5u) { "unlikely" } else { "likely" });
|
|
|
|
let a = [mut 1, 2, 3];
|
|
shuffle(r, a);
|
|
log(error, a);
|
|
|
|
let i = 0u;
|
|
let v = [
|
|
{weight:1u, item:"low"},
|
|
{weight:8u, item:"middle"},
|
|
{weight:1u, item:"high"}
|
|
];
|
|
let w = weighted_vec(v);
|
|
|
|
while i < 1000u {
|
|
log(error, "Immed: " + weighted_choice(r, v));
|
|
log(error, "Fast: " + choice(r, w));
|
|
i += 1u;
|
|
}
|
|
}
|