Implement Roots for BigInt and BigUint

Cargo.toml

@@ -31,7 +31,7 @@ name = "shootout-pidigits"
 [dependencies]
 
 [dependencies.num-integer]
-version = "0.1.38"
+version = "0.1.39"
 default-features = false
 
 [dependencies.num-traits]

benches/bigint.rs

@@ -4,13 +4,14 @@
 extern crate test;
 extern crate num_bigint;
 extern crate num_traits;
+extern crate num_integer;
 extern crate rand;
 
 use std::mem::replace;
 use test::Bencher;
 use num_bigint::{BigInt, BigUint, RandBigInt};
 use num_traits::{Zero, One, FromPrimitive, Num};
-use rand::{SeedableRng, StdRng};
+use rand::{SeedableRng, StdRng, Rng};
 
 fn get_rng() -> StdRng {
     let mut seed = [0; 32];
@@ -342,3 +343,32 @@ fn modpow_even(b: &mut Bencher) {
 
     b.iter(|| base.modpow(&e, &m));
 }
+
+#[bench]
+fn roots_sqrt(b: &mut Bencher) {
+    let mut rng = get_rng();
+    let x = rng.gen_biguint(2048);
+
+    b.iter(|| x.sqrt());
+}
+
+#[bench]
+fn roots_cbrt(b: &mut Bencher) {
+    let mut rng = get_rng();
+    let x = rng.gen_biguint(2048);
+
+    b.iter(|| x.cbrt());
+}
+
+#[bench]
+fn roots_nth(b: &mut Bencher) {
+    let mut rng = get_rng();
+    let x = rng.gen_biguint(2048);
+    // Although n is u32, here we limit it to the set of u8 values since it
+    // hugely impacts the performance of nth_root due to exponentiation to
+    // the power of n-1. Using very large values for n is also not very realistic,
+    // and any n > x's bit size produces 1 as a result anyway.
+    let n: u8 = rng.gen();
+
+    b.iter(|| { x.nth_root(n as u32) });
+}

src/bigint.rs

@@ -16,7 +16,7 @@ use std::iter::{Product, Sum};
 #[cfg(feature = "serde")]
 use serde;
 
-use integer::Integer;
+use integer::{Integer, Roots};
 use traits::{ToPrimitive, FromPrimitive, Num, CheckedAdd, CheckedSub,
              CheckedMul, CheckedDiv, Signed, Zero, One};
 
@@ -1802,6 +1802,15 @@ impl Integer for BigInt {
     }
 }
 
+impl Roots for BigInt {
+    fn nth_root(&self, n: u32) -> Self {
+        assert!(!(self.is_negative() && n.is_even()),
+                "n-th root is undefined for number (n={})", n);
+
+        BigInt::from_biguint(self.sign, self.data.nth_root(n))
+    }
+}
+
 impl ToPrimitive for BigInt {
     #[inline]
     fn to_i64(&self) -> Option<i64> {
@@ -2538,6 +2547,25 @@ impl BigInt {
         };
         BigInt::from_biguint(sign, mag)
     }
+
+    /// Returns the truncated principal square root of `self` --
+    /// see [Roots::sqrt](Roots::sqrt).
+    // struct.BigInt.html#trait.Roots
+    pub fn sqrt(&self) -> Self {
+        Roots::sqrt(self)
+    }
+
+    /// Returns the truncated principal cube root of `self` --
+    /// see [Roots::cbrt](Roots::cbrt).
+    pub fn cbrt(&self) -> Self {
+        Roots::cbrt(self)
+    }
+
+    /// Returns the truncated principal `n`th root of `self` --
+    /// See [Roots::nth_root](Roots::nth_root).
+    pub fn nth_root(&self, n: u32) -> Self {
+        Roots::nth_root(self, n)
+    }
 }
 
 impl_sum_iter_type!(BigInt);

src/biguint.rs

@@ -17,9 +17,9 @@ use std::ascii::AsciiExt;
 #[cfg(feature = "serde")]
 use serde;
 
-use integer::Integer;
+use integer::{Integer, Roots};
 use traits::{ToPrimitive, FromPrimitive, Float, Num, Unsigned, CheckedAdd, CheckedSub, CheckedMul,
-             CheckedDiv, Zero, One};
+             CheckedDiv, Zero, One, pow};
 
 use big_digit::{self, BigDigit, DoubleBigDigit};
 
@@ -1026,6 +1026,52 @@ impl Integer for BigUint {
     }
 }
 
+impl Roots for BigUint {
+    fn nth_root(&self, n: u32) -> Self {
+        assert!(n > 0, "n must be at least 1");
+
+        let one = BigUint::one();
+
+        // Trivial cases
+        if self.is_zero() {
+            return BigUint::zero();
+        }
+
+        if self.is_one() {
+            return one;
+        }
+
+        let n = n as usize;
+        let n_min_1 = (n as usize) - 1;
+
+        // Newton's method to compute the nth root of an integer.
+        //
+        // Reference:
+        // Brent & Zimmermann, Modern Computer Arithmetic, v0.5.9, Algorithm 1.14
+        //
+        // Set initial guess to something definitely >= floor(nth_root of self)
+        // but as low as possible to speed up convergence.
+        let bit_len = self.len() * big_digit::BITS;
+        let guess = one << (bit_len/n + 1);
+
+        let mut u = guess;
+        let mut s: BigUint;
+
+        loop {
+            s = u;
+            let q = self / pow(s.clone(), n_min_1);
+            let t: BigUint = n_min_1 * &s + q;
+
+            // Compute the candidate value for next iteration
+            u = t / n;
+
+            if u >= s { break; }
+        }
+
+        s
+    }
+}
+
 fn high_bits_to_u64(v: &BigUint) -> u64 {
     match v.data.len() {
         0   => 0,
@@ -1749,6 +1795,25 @@ impl BigUint {
         }
         acc
     }
+
+    /// Returns the truncated principal square root of `self` --
+    /// see [Roots::sqrt](Roots::sqrt).
+    // struct.BigInt.html#trait.Roots
+    pub fn sqrt(&self) -> Self {
+        Roots::sqrt(self)
+    }
+
+    /// Returns the truncated principal cube root of `self` --
+    /// see [Roots::cbrt](Roots::cbrt).
+    pub fn cbrt(&self) -> Self {
+        Roots::cbrt(self)
+    }
+
+    /// Returns the truncated principal `n`th root of `self` --
+    /// See [Roots::nth_root](Roots::nth_root).
+    pub fn nth_root(&self, n: u32) -> Self {
+        Roots::nth_root(self, n)
+    }
 }
 
 /// Returns the number of least-significant bits that are zero,

tests/roots.rs

@@ -0,0 +1,84 @@
+extern crate num_bigint;
+extern crate num_integer;
+extern crate num_traits;
+
+mod biguint {
+    use num_bigint::BigUint;
+    use num_traits::FromPrimitive;
+    use std::str::FromStr;
+
+    fn check(x: i32, n: u32, expected: i32) {
+        let big_x: BigUint = FromPrimitive::from_i32(x).unwrap();
+        let big_expected: BigUint = FromPrimitive::from_i32(expected).unwrap();
+
+        assert_eq!(big_x.nth_root(n), big_expected);
+    }
+
+    #[test]
+    fn test_sqrt() {
+        check(99, 2, 9);
+        check(100, 2, 10);
+        check(120, 2, 10);
+    }
+
+    #[test]
+    fn test_cbrt() {
+        check(8, 3, 2);
+        check(26, 3, 2);
+    }
+
+    #[test]
+    fn test_nth_root() {
+        check(0, 1, 0);
+        check(10, 1, 10);
+        check(100, 4, 3);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_nth_root_n_is_zero() {
+        check(4, 0, 0);
+    }
+
+    #[test]
+    fn test_nth_root_big() {
+        let x: BigUint = FromStr::from_str("123_456_789").unwrap();
+        let expected : BigUint = FromPrimitive::from_i32(6).unwrap();
+
+        assert_eq!(x.nth_root(10), expected);
+    }
+}
+
+mod bigint {
+    use num_bigint::BigInt;
+    use num_traits::FromPrimitive;
+
+    fn check(x: i32, n: u32, expected: i32) {
+        let big_x: BigInt = FromPrimitive::from_i32(x).unwrap();
+        let big_expected: BigInt = FromPrimitive::from_i32(expected).unwrap();
+
+        assert_eq!(big_x.nth_root(n), big_expected);
+    }
+
+    #[test]
+    fn test_nth_root() {
+        check(-100, 3, -4);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_nth_root_x_neg_n_even() {
+        check(-100, 4, 0);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_sqrt_x_neg() {
+        check(-4, 2, -2);
+    }
+
+    #[test]
+    fn test_cbrt() {
+        check(-8, 3, -2);
+    }
+}