|
| 1 | +use std; |
| 2 | +import std::rand; |
| 3 | +import std::vec; |
| 4 | + |
| 5 | +// random uint less than n |
| 6 | +fn under(r : rand::rng, n : uint) -> uint { assert n != 0u; r.next() as uint % n } |
| 7 | + |
| 8 | +// random choice from a vec |
| 9 | +fn choice<T>(r : rand::rng, v : [T]) -> T { assert vec::len(v) != 0u; v[under(r, vec::len(v))] } |
| 10 | + |
| 11 | +// 1 in n chance of being true |
| 12 | +fn unlikely(r : rand::rng, n : uint) -> bool { under(r, n) == 0u } |
| 13 | + |
| 14 | +// shuffle a vec in place |
| 15 | +fn shuffle<@T>(r : rand::rng, &v : [mutable T]) { |
| 16 | + let i = vec::len(v); |
| 17 | + while i >= 2u { |
| 18 | + // Loop invariant: elements with index >= i have been locked in place. |
| 19 | + i -= 1u; |
| 20 | + vec::swap(v, i, under(r, i + 1u)); // Lock element i in place. |
| 21 | + } |
| 22 | +} |
| 23 | + |
| 24 | +// create a shuffled copy of a vec |
| 25 | +fn shuffled<@T>(r : rand::rng, v : [T]) -> [T] { |
| 26 | + let w = vec::to_mut(v); |
| 27 | + shuffle(r, w); |
| 28 | + vec::from_mut(w) // Shouldn't this happen automatically? |
| 29 | +} |
| 30 | + |
| 31 | +// sample from a population without replacement |
| 32 | +//fn sample<T>(r : rand::rng, pop : [T], k : uint) -> [T] { fail } |
| 33 | + |
| 34 | +// Two ways to make a weighted choice. |
| 35 | +// * weighted_choice is O(number of choices) time |
| 36 | +// * weighted_vec is O(total weight) space |
| 37 | +type weighted<T> = { weight: uint, item: T }; |
| 38 | +fn weighted_choice<@T>(r : rand::rng, v : [weighted<T>]) -> T { |
| 39 | + assert vec::len(v) != 0u; |
| 40 | + let total = 0u; |
| 41 | + for {weight: weight, item: _} in v { |
| 42 | + total += weight; |
| 43 | + } |
| 44 | + assert total >= 0u; |
| 45 | + let chosen = under(r, total); |
| 46 | + let so_far = 0u; |
| 47 | + for {weight: weight, item: item} in v { |
| 48 | + so_far += weight; |
| 49 | + if so_far > chosen { |
| 50 | + ret item; |
| 51 | + } |
| 52 | + } |
| 53 | + std::util::unreachable(); |
| 54 | +} |
| 55 | + |
| 56 | +fn weighted_vec<T>(v : [weighted<T>]) -> [T] { |
| 57 | + let r = []; |
| 58 | + for {weight: weight, item: item} in v { |
| 59 | + let i = 0u; |
| 60 | + while i < weight { |
| 61 | + r += [item]; |
| 62 | + i += 1u; |
| 63 | + } |
| 64 | + } |
| 65 | + r |
| 66 | +} |
| 67 | + |
| 68 | +fn main() |
| 69 | +{ |
| 70 | + let r = rand::mk_rng(); |
| 71 | + |
| 72 | + log_err under(r, 5u); |
| 73 | + log_err choice(r, [10, 20, 30]); |
| 74 | + log_err if unlikely(r, 5u) { "unlikely" } else { "likely" }; |
| 75 | + |
| 76 | + let a = [mutable 1, 2, 3]; |
| 77 | + shuffle(r, a); |
| 78 | + log_err a; |
| 79 | + |
| 80 | + let i = 0u; |
| 81 | + let v = [ |
| 82 | + {weight:1u, item:"low"}, |
| 83 | + {weight:8u, item:"middle"}, |
| 84 | + {weight:1u, item:"high"} |
| 85 | + ]; |
| 86 | + let w = weighted_vec(v); |
| 87 | + |
| 88 | + while i < 1000u { |
| 89 | + log_err "Immed: " + weighted_choice(r, v); |
| 90 | + log_err "Fast: " + choice(r, w); |
| 91 | + i += 1u; |
| 92 | + } |
| 93 | +} |
0 commit comments