diff --git a/src/distributions/binomial.rs b/src/distributions/binomial.rs index 907155b24ed..2bad03b78bd 100644 --- a/src/distributions/binomial.rs +++ b/src/distributions/binomial.rs @@ -47,6 +47,31 @@ impl Binomial { } } +/// Raise a `base` to the power of `exp`, using exponentiation by squaring. +/// +/// This implementation is based on the one in the `num_traits` crate. It is +/// slightly modified to accept `u64` exponents. +fn pow(mut base: f64, mut exp: u64) -> f64 { + if exp == 0 { + return 1.; + } + + while exp & 1 == 0 { + base *= base; + exp >>= 1; + } + + let mut acc = base; + while exp > 1 { + exp >>= 1; + base *= base; + if exp & 1 == 1 { + acc *= base; + } + } + acc +} + impl Distribution for Binomial { fn sample(&self, rng: &mut R) -> u64 { // Handle these values directly. @@ -77,9 +102,7 @@ impl Distribution for Binomial { let q = 1. - p; let s = p / q; let a = ((self.n + 1) as f64) * s; - let mut r = q.powf(self.n as f64); - // FIXME: Using integer exponentiation should be faster, but it is not available for - // `u64` exponents, neither in `std` or `num`. + let mut r = pow(q, self.n); let mut u: f64 = rng.gen(); let mut x = 0; while u > r as f64 {