結果
| 問題 | No.5007 Steiner Space Travel |
| ユーザー |
 ntk-ta01
|
| 提出日時 | 2022-07-30 17:34:05 |
| 言語 | Rust (1.77.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 962 ms / 1,000 ms |
| コード長 | 26,780 bytes |
| コンパイル時間 | 1,390 ms |
| 実行使用メモリ | 6,952 KB |
| スコア | 8,810,615 |
| 最終ジャッジ日時 | 2022-07-30 17:34:38 |
| 合計ジャッジ時間 | 32,982 ms |
|
ジャッジサーバーID （参考情報） |
judge10 / judge15 |
| 純コード判定しない問題か言語 |
（要ログイン）
テストケース
テストケース表示| 入力 | 結果 | 実行時間 実行使用メモリ |
|---|---|---|
| testcase_00 | AC | 961 ms
5,156 KB |
| testcase_01 | AC | 961 ms
4,900 KB |
| testcase_02 | AC | 961 ms
4,904 KB |
| testcase_03 | AC | 961 ms
4,904 KB |
| testcase_04 | AC | 961 ms
4,900 KB |
| testcase_05 | AC | 961 ms
4,904 KB |
| testcase_06 | AC | 962 ms
4,900 KB |
| testcase_07 | AC | 962 ms
4,900 KB |
| testcase_08 | AC | 960 ms
4,900 KB |
| testcase_09 | AC | 962 ms
4,904 KB |
| testcase_10 | AC | 961 ms
4,900 KB |
| testcase_11 | AC | 961 ms
4,904 KB |
| testcase_12 | AC | 962 ms
4,900 KB |
| testcase_13 | AC | 961 ms
6,952 KB |
| testcase_14 | AC | 962 ms
6,948 KB |
| testcase_15 | AC | 961 ms
4,904 KB |
| testcase_16 | AC | 961 ms
4,900 KB |
| testcase_17 | AC | 962 ms
6,948 KB |
| testcase_18 | AC | 961 ms
4,908 KB |
| testcase_19 | AC | 961 ms
4,900 KB |
| testcase_20 | AC | 962 ms
4,900 KB |
| testcase_21 | AC | 962 ms
4,904 KB |
| testcase_22 | AC | 961 ms
6,948 KB |
| testcase_23 | AC | 961 ms
4,904 KB |
| testcase_24 | AC | 962 ms
4,900 KB |
| testcase_25 | AC | 961 ms
4,908 KB |
| testcase_26 | AC | 962 ms
6,952 KB |
| testcase_27 | AC | 961 ms
6,948 KB |
| testcase_28 | AC | 962 ms
4,900 KB |
| testcase_29 | AC | 962 ms
4,904 KB |
ソースコード
// https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
macro_rules! input {
(source = $s:expr, $($r:tt)*) => {
let mut iter = $s.split_whitespace();
input_inner!{iter, $($r)*}
};
($($r:tt)*) => {
let s = {
use std::io::Read;
let mut s = String::new();
std::io::stdin().read_to_string(&mut s).unwrap();
s
};
let mut iter = s.split_whitespace();
input_inner!{iter, $($r)*}
};
}
macro_rules! input_inner {
($iter:expr) => {};
($iter:expr, ) => {};
($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
let $var = read_value!($iter, $t);
input_inner!{$iter $($r)*}
};
}
macro_rules! read_value {
($iter:expr, ( $($t:tt),* )) => {
( $(read_value!($iter, $t)),* )
};
($iter:expr, [ $t:tt ; $len:expr ]) => {
(0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
};
($iter:expr, chars) => {
read_value!($iter, String).chars().collect::<Vec<char>>()
};
($iter:expr, usize1) => {
read_value!($iter, usize) - 1
};
($iter:expr, $t:ty) => {
$iter.next().unwrap().parse::<$t>().expect("Parse error")
};
}
use ntk_rand::Xorshift;
const TIMELIMIT: f64 = 0.24;
fn main() {
let input = Input::parse_input();
let mut solutions: Vec<Solution> = vec![];
for i in 1..5 {
let mut timer = Timer::new();
let mut rng = Xorshift::seed_from_u64(i * 9050230);
let mut solution = Solution::new(&input);
if 1 < i {
solution.stations = solutions[i - 2].stations.clone();
let mut visits = vec![];
for &v in solutions[i - 2].visits.iter().filter(|v| v.0 == 1) {
// 惑星なら追加
visits.push(v);
}
solution.visits = visits;
solution.visit_time = solutions[i - 2].visit_time.clone();
}
annealing(&input, &mut solution, &mut timer, &mut rng);
solutions.push(solution)
}
let mut best_solution_i = 0;
let mut best_score = 0;
for (i, solution) in solutions.iter().enumerate() {
if best_score < solution.score {
best_score = solution.score;
best_solution_i = i;
}
}
println!("{}", solutions[best_solution_i]);
eprintln!("{}", best_score);
}
fn annealing(input: &Input, solution: &mut Solution, timer: &mut Timer, rng: &mut Xorshift<usize>) {
const T0: f64 = 100.0;
const T1: f64 = 0.00001;
let mut temp = T0;
let mut prob;
solution.compute_score(input);
let mut best_solution = solution.clone();
let mut count = 0;
loop {
if count >= 100 {
let passed = timer.get_time() / TIMELIMIT;
if passed >= 1.0 {
break;
}
temp = T0.powf(1.0 - passed) * T1.powf(passed);
count = 0;
}
count += 1;
let mut new_solution = solution.clone();
// 近傍解作成
let neigh_count = 6;
match rng.gen_range(0..neigh_count) {
0 => {
// 訪問順を変える swap
let v1 = rng.gen_range(1..new_solution.visits.len() - 1);
let v2 = rng.gen_range(1..new_solution.visits.len() - 1);
new_solution.visits.swap(v1, v2);
}
1 => {
// 2-opt
let mut i = rng.gen_range(1..new_solution.visits.len() - 1);
let mut j = rng.gen_range(1..new_solution.visits.len() - 1);
if i == j {
continue;
}
if i > j {
std::mem::swap(&mut i, &mut j);
}
new_solution.visits[i..=j].reverse();
}
2 => {
// 惑星の経由地を追加する
let planet = rng.gen_range(0..input.n);
let i = rng.gen_range(1..new_solution.visits.len());
new_solution.visits.insert(i, (1, planet));
new_solution.visit_time[planet] += 1;
}
3 => {
// 経由地を削除する 惑星 or ステーション
// TODO: ステーション → ステーションの訪問になるようなら上手くやる やれ
let i = rng.gen_range(1..new_solution.visits.len() - 1);
let astronomy = new_solution.visits[i];
if astronomy.0 == 1 && new_solution.visit_time[astronomy.1] == 1 {
continue;
}
new_solution.visits.remove(i);
if astronomy.0 == 1 {
new_solution.visit_time[astronomy.1] -= 1;
}
}
4 => {
// ステーションを訪問させる insert
// ランダムに移動
let station = rng.gen_range(0..input.m);
let i = rng.gen_range(1..new_solution.visits.len());
new_solution.visits.insert(i, (2, station));
// 最良移動戦略で書く あんまりよくならそう
// let mut best_i = 0;
// let mut best_station = 0;
// let mut max_diff = 0;
// for i in 1..new_solution.visits.len() {
// let prev_planet = new_solution.visits[i - 1];
// let next_planet = new_solution.visits[i];
// if prev_planet.0 == 2 || next_planet.0 == 2 {
// // ステーション → ステーションの訪問はなし
// continue;
// }
// let p1 = input.planets[prev_planet.1];
// let p2 = input.planets[next_planet.1];
// let prev_dist = ALPHA * ALPHA * compute_squared_distance(p1, p2);
// for station in 0..input.m {
// let s = new_solution.stations[station];
// let next_dist = ALPHA * compute_squared_distance(s, p1)
// + ALPHA * compute_squared_distance(s, p2);
// if next_dist >= prev_dist {
// continue;
// }
// let diff = prev_dist - next_dist;
// if max_diff < diff {
// best_i = i;
// best_station = station;
// max_diff = diff;
// }
// }
// }
// new_solution.visits.insert(best_i, (2, best_station));
}
5 => {
// ステーションを移動させる move_station
let i = rng.gen_range(0..input.m);
new_solution.stations[i].0 += if rng.gen_bool(0.5) { 1 } else { -1 };
new_solution.stations[i].1 += if rng.gen_bool(0.5) { 1 } else { -1 };
}
_ => unreachable!(),
}
// 近傍解作成ここまで
new_solution.compute_score(input);
prob = f64::exp((new_solution.score - solution.score) as f64 / temp);
if solution.score < new_solution.score || rng.gen_bool(prob) {
*solution = new_solution;
}
if best_solution.score < solution.score {
best_solution = solution.clone();
}
}
*solution = best_solution;
}
const ALPHA: i64 = 5;
#[derive(Debug, Clone)]
struct Input {
n: usize,
m: usize,
planets: Vec<(i64, i64)>,
}
impl Input {
fn parse_input() -> Self {
input! {
n: usize,
m: usize,
planets: [(i64, i64); n],
}
Input { n, m, planets }
}
}
#[derive(Debug, Clone)]
struct Solution {
stations: Vec<(i64, i64)>,
visits: Vec<(usize, usize)>, // lenはinput.n + 1以上, (usize, Usize1)の気持ち
visit_time: Vec<usize>, // 各惑星の訪問回数 1以上
score: i64,
}
impl std::fmt::Display for Solution {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for (c, d) in self.stations.iter() {
writeln!(f, "{} {}", c, d)?;
}
writeln!(f, "{}", self.visits.len())?;
for (t, r) in self.visits.iter() {
writeln!(f, "{} {}", t, r + 1)?;
}
Ok(())
}
}
fn compute_squared_distance(p1: (i64, i64), p2: (i64, i64)) -> i64 {
let dx = p1.0 - p2.0;
let dy = p1.1 - p2.1;
dx * dx + dy * dy
}
impl Solution {
fn new(input: &Input) -> Self {
let stations = vec![
(200, 400),
(400, 400),
(600, 400),
(800, 400),
(200, 600),
(400, 600),
(600, 600),
(800, 600),
];
let mut visits: Vec<(usize, usize)> = (0..input.n).into_iter().map(|r| (1, r)).collect();
visits.push((1, 0));
let visit_time = vec![1; input.n];
Solution {
stations,
visits,
visit_time,
score: 0,
}
}
fn compute_score(&mut self, input: &Input) {
let mut score = 0;
for i in 0..self.visits.len() - 1 {
let prev_type = self.visits[i].0;
let prev_point = if prev_type == 1 {
input.planets[self.visits[i].1]
} else {
self.stations[self.visits[i].1]
};
let next_type = self.visits[i + 1].0;
let next_point = if next_type == 1 {
input.planets[self.visits[i + 1].1]
} else {
self.stations[self.visits[i + 1].1]
};
let base_energy = compute_squared_distance(prev_point, next_point);
let prev_multiplier = if prev_type == 1 { ALPHA } else { 1 };
let next_multiplier = if next_type == 1 { ALPHA } else { 1 };
score += base_energy * prev_multiplier * next_multiplier;
}
let score = (1e9 / (1000.0 + f64::sqrt(score as f64))).round() as i64;
self.score = score;
}
}
pub fn get_time() -> f64 {
let t = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap();
t.as_secs() as f64 + t.subsec_nanos() as f64 * 1e-9
}
struct Timer {
start_time: f64,
}
impl Timer {
fn new() -> Timer {
Timer {
start_time: get_time(),
}
}
fn get_time(&self) -> f64 {
get_time() - self.start_time
}
}
#[allow(dead_code)]
mod ntk_rand {
pub trait Distribution<T, TX> {
/// Generate a random value of `T`, using `rng` as the source of randomness.
fn sample(&self, rng: &mut Xorshift<TX>) -> T;
}
pub struct Xorshift<T> {
seed: T,
}
impl Xorshift<usize> {
pub fn new() -> Self {
Xorshift {
seed: 0x139408dcbbf7a44,
}
}
pub fn seed_from_u64(seed: usize) -> Xorshift<usize> {
Xorshift { seed }
}
fn gen(&mut self) -> usize {
self.seed = self.seed ^ (self.seed << 13);
self.seed = self.seed ^ (self.seed >> 7);
self.seed = self.seed ^ (self.seed << 17);
self.seed
}
pub fn rand(&mut self) -> usize {
self.gen()
}
fn sample<T, D: Distribution<T, usize>>(&mut self, distr: D) -> T {
distr.sample(self)
}
pub fn gen_bool(&mut self, p: f64) -> bool {
let d = Bernoulli::new(p).unwrap();
self.sample(d)
}
pub fn gen_range<T, R>(&mut self, range: R) -> T
where
T: SampleUniform,
R: SampleRange<T>,
{
range.sample_single(self)
}
}
pub struct Uniform<X: SampleUniform>(X::Sampler);
impl<X: SampleUniform> Uniform<X> {
pub fn new<B1, B2>(low: B1, high: B2) -> Uniform<X>
where
B1: SampleBorrow<X> + Sized,
B2: SampleBorrow<X> + Sized,
{
Uniform(X::Sampler::new(low, high))
}
pub fn new_inclusive<B1, B2>(low: B1, high: B2) -> Uniform<X>
where
B1: SampleBorrow<X> + Sized,
B2: SampleBorrow<X> + Sized,
{
Uniform(X::Sampler::new_inclusive(low, high))
}
}
impl<X: SampleUniform> Distribution<X, usize> for Uniform<X> {
fn sample(&self, rng: &mut Xorshift<usize>) -> X {
self.0.sample(rng)
}
}
pub trait UniformSampler: Sized {
/// The type sampled by this implementation.
type X;
fn new<B1, B2>(low: B1, high: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized;
fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized;
/// Sample a value.
fn sample(&self, rng: &mut Xorshift<usize>) -> Self::X;
fn sample_single<B1, B2>(low: B1, high: B2, rng: &mut Xorshift<usize>) -> Self::X
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let uniform: Self = UniformSampler::new(low, high);
uniform.sample(rng)
}
fn sample_single_inclusive<B1, B2>(low: B1, high: B2, rng: &mut Xorshift<usize>) -> Self::X
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let uniform: Self = UniformSampler::new_inclusive(low, high);
uniform.sample(rng)
}
}
impl<X: SampleUniform> From<core::ops::Range<X>> for Uniform<X> {
fn from(r: core::ops::Range<X>) -> Uniform<X> {
Uniform::new(r.start, r.end)
}
}
impl<X: SampleUniform> From<core::ops::RangeInclusive<X>> for Uniform<X> {
fn from(r: core::ops::RangeInclusive<X>) -> Uniform<X> {
Uniform::new_inclusive(r.start(), r.end())
}
}
pub trait SampleUniform: Sized {
/// The `UniformSampler` implementation supporting type `X`.
type Sampler: UniformSampler<X = Self>;
}
pub trait SampleBorrow<Borrowed> {
/// Immutably borrows from an owned value. See [`Borrow::borrow`]
///
/// [`Borrow::borrow`]: std::borrow::Borrow::borrow
fn borrow(&self) -> &Borrowed;
}
impl<Borrowed> SampleBorrow<Borrowed> for Borrowed
where
Borrowed: SampleUniform,
{
#[inline(always)]
fn borrow(&self) -> &Borrowed {
self
}
}
impl<'a, Borrowed> SampleBorrow<Borrowed> for &'a Borrowed
where
Borrowed: SampleUniform,
{
#[inline(always)]
fn borrow(&self) -> &Borrowed {
*self
}
}
/// Range that supports generating a single sample efficiently.
///
/// Any type implementing this trait can be used to specify the sampled range
/// for `Rng::gen_range`.
pub trait SampleRange<T> {
/// Generate a sample from the given range.
fn sample_single(self, rng: &mut Xorshift<usize>) -> T;
/// Check whether the range is empty.
fn is_empty(&self) -> bool;
}
impl<T: SampleUniform + PartialOrd> SampleRange<T> for core::ops::Range<T> {
#[inline]
fn sample_single(self, rng: &mut Xorshift<usize>) -> T {
T::Sampler::sample_single(self.start, self.end, rng)
}
#[inline]
fn is_empty(&self) -> bool {
// !(self.start >= self.end)
self.start < self.end
}
}
impl<T: SampleUniform + PartialOrd> SampleRange<T> for core::ops::RangeInclusive<T> {
#[inline]
fn sample_single(self, rng: &mut Xorshift<usize>) -> T {
T::Sampler::sample_single_inclusive(self.start(), self.end(), rng)
}
#[inline]
fn is_empty(&self) -> bool {
// !(self.start() <= self.end())
self.start() > self.end()
}
}
pub trait WideningMultiply<RHS = Self> {
type Output;
fn wmul(self, x: RHS) -> Self::Output;
}
macro_rules! wmul_impl {
($ty:ty, $wide:ty, $shift:expr) => {
impl WideningMultiply for $ty {
type Output = ($ty, $ty);
#[inline(always)]
fn wmul(self, x: $ty) -> Self::Output {
let tmp = (self as $wide) * (x as $wide);
((tmp >> $shift) as $ty, tmp as $ty)
}
}
};
// simd bulk implementation
($(($ty:ident, $wide:ident),)+, $shift:expr) => {
$(
impl WideningMultiply for $ty {
type Output = ($ty, $ty);
#[inline(always)]
fn wmul(self, x: $ty) -> Self::Output {
// For supported vectors, this should compile to a couple
// supported multiply & swizzle instructions (no actual
// casting).
// TODO: optimize
let y: $wide = self.cast();
let x: $wide = x.cast();
let tmp = y * x;
let hi: $ty = (tmp >> $shift).cast();
let lo: $ty = tmp.cast();
(hi, lo)
}
}
)+
};
}
wmul_impl! { u8, u16, 8 }
wmul_impl! { u16, u32, 16 }
wmul_impl! { u32, u64, 32 }
wmul_impl! { u64, u128, 64 }
macro_rules! wmul_impl_large {
($ty:ty, $half:expr) => {
impl WideningMultiply for $ty {
type Output = ($ty, $ty);
#[inline(always)]
fn wmul(self, b: $ty) -> Self::Output {
const LOWER_MASK: $ty = !0 >> $half;
let mut low = (self & LOWER_MASK).wrapping_mul(b & LOWER_MASK);
let mut t = low >> $half;
low &= LOWER_MASK;
t += (self >> $half).wrapping_mul(b & LOWER_MASK);
low += (t & LOWER_MASK) << $half;
let mut high = t >> $half;
t = low >> $half;
low &= LOWER_MASK;
t += (b >> $half).wrapping_mul(self & LOWER_MASK);
low += (t & LOWER_MASK) << $half;
high += t >> $half;
high += (self >> $half).wrapping_mul(b >> $half);
(high, low)
}
}
};
// simd bulk implementation
(($($ty:ty,)+) $scalar:ty, $half:expr) => {
$(
impl WideningMultiply for $ty {
type Output = ($ty, $ty);
#[inline(always)]
fn wmul(self, b: $ty) -> Self::Output {
// needs wrapping multiplication
const LOWER_MASK: $scalar = !0 >> $half;
let mut low = (self & LOWER_MASK) * (b & LOWER_MASK);
let mut t = low >> $half;
low &= LOWER_MASK;
t += (self >> $half) * (b & LOWER_MASK);
low += (t & LOWER_MASK) << $half;
let mut high = t >> $half;
t = low >> $half;
low &= LOWER_MASK;
t += (b >> $half) * (self & LOWER_MASK);
low += (t & LOWER_MASK) << $half;
high += t >> $half;
high += (self >> $half) * (b >> $half);
(high, low)
}
}
)+
};
}
wmul_impl_large! { u128, 64 }
macro_rules! wmul_impl_usize {
($ty:ty) => {
impl WideningMultiply for usize {
type Output = (usize, usize);
#[inline(always)]
fn wmul(self, x: usize) -> Self::Output {
let (high, low) = (self as $ty).wmul(x as $ty);
(high as usize, low as usize)
}
}
};
}
#[cfg(target_pointer_width = "16")]
wmul_impl_usize! { u16 }
#[cfg(target_pointer_width = "32")]
wmul_impl_usize! { u32 }
#[cfg(target_pointer_width = "64")]
wmul_impl_usize! { u64 }
pub struct UniformInt<X> {
low: X,
range: X,
z: X,
}
macro_rules! uniform_int_impl {
($ty:ty, $unsigned:ident, $u_large:ident) => {
impl SampleUniform for $ty {
type Sampler = UniformInt<$ty>;
}
impl UniformSampler for UniformInt<$ty> {
type X = $ty;
#[inline]
fn new<B1, B2>(low_b: B1, high_b: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let low = *low_b.borrow();
let high = *high_b.borrow();
assert!(low < high, "Uniform::new called with `low >= high`");
UniformSampler::new_inclusive(low, high - 1)
}
#[inline]
fn new_inclusive<B1, B2>(low_b: B1, high_b: B2) -> Self
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let low = *low_b.borrow();
let high = *high_b.borrow();
assert!(
low <= high,
"Uniform::new_inclusive called with `low > high`"
);
let unsigned_max = core::$u_large::MAX;
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned;
let ints_to_reject = if range > 0 {
let range = $u_large::from(range);
(unsigned_max - range + 1) % range
} else {
0
};
UniformInt {
low,
range: range as $ty,
z: ints_to_reject as $unsigned as $ty,
}
}
#[inline]
fn sample(&self, rng: &mut Xorshift<$u_large>) -> Self::X {
let range = self.range as $unsigned as $u_large;
if range > 0 {
let unsigned_max = core::$u_large::MAX;
let zone = unsigned_max - (self.z as $unsigned as $u_large);
loop {
let v: $u_large = rng.gen();
let (hi, lo) = v.wmul(range);
if lo <= zone {
return self.low.wrapping_add(hi as $ty);
}
}
} else {
rng.gen()
}
}
#[inline]
fn sample_single<B1, B2>(
low_b: B1,
high_b: B2,
rng: &mut Xorshift<$u_large>,
) -> Self::X
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let low = *low_b.borrow();
let high = *high_b.borrow();
assert!(low < high, "UniformSampler::sample_single: low >= high");
Self::sample_single_inclusive(low, high - 1, rng)
}
#[inline]
fn sample_single_inclusive<B1, B2>(
low_b: B1,
high_b: B2,
rng: &mut Xorshift<$u_large>,
) -> Self::X
where
B1: SampleBorrow<Self::X> + Sized,
B2: SampleBorrow<Self::X> + Sized,
{
let low = *low_b.borrow();
let high = *high_b.borrow();
assert!(
low <= high,
"UniformSampler::sample_single_inclusive: low > high"
);
let range = high.wrapping_sub(low).wrapping_add(1) as $unsigned as $u_large;
if range == 0 {
return rng.gen();
}
let zone = if core::$unsigned::MAX <= core::u16::MAX as $unsigned {
let unsigned_max: $u_large = core::$u_large::MAX;
let ints_to_reject = (unsigned_max - range + 1) % range;
unsigned_max - ints_to_reject
} else {
(range << range.leading_zeros()).wrapping_sub(1)
};
loop {
let v: $u_large = rng.gen();
let (hi, lo) = v.wmul(range);
if lo <= zone {
return low.wrapping_add(hi as $ty);
}
}
}
}
};
}
uniform_int_impl! { usize, usize, usize }
pub struct Bernoulli {
p_int: u64,
}
const ALWAYS_TRUE: u64 = u64::max_value();
const SCALE: f64 = 2.0 * (1u64 << 63) as f64;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BernoulliError {
InvalidProbability,
}
impl Bernoulli {
#[inline]
pub fn new(p: f64) -> Result<Bernoulli, BernoulliError> {
if !(0.0..1.0).contains(&p) {
if p == 1.0 {
return Ok(Bernoulli { p_int: ALWAYS_TRUE });
}
return Err(BernoulliError::InvalidProbability);
}
Ok(Bernoulli {
p_int: (p * SCALE) as u64,
})
}
}
impl Distribution<bool, usize> for Bernoulli {
#[inline]
fn sample(&self, rng: &mut Xorshift<usize>) -> bool {
if self.p_int == ALWAYS_TRUE {
return true;
}
let v: u64 = rng.gen() as u64;
v < self.p_int
}
}
}