浏览代码

Auto merge of #272 - vks:binomial-coeffs, r=cuviper

Implement an iterator over the binomial coefficients

I'm not very happy with the excessive cloning, but to fix it the bounds on the type parameters would have to be excessive. We probably need something like [this](https://github.com/vks/discrete-log/blob/master/src/main.rs#L90) in `num-traits`.
Homu 8 年之前
父节点
当前提交
f63c933737
共有 1 个文件被更改,包括 121 次插入1 次删除
  1. 121 1
      integer/src/lib.rs

+ 121 - 1
integer/src/lib.rs

@@ -667,14 +667,91 @@ impl_integer_for_usize!(u32, test_integer_u32);
 impl_integer_for_usize!(u64, test_integer_u64);
 impl_integer_for_usize!(usize, test_integer_usize);
 
+/// An iterator over binomial coefficients.
+pub struct IterBinomial<T> {
+    a: T,
+    n: T,
+    k: T,
+}
+
+impl<T> IterBinomial<T>
+    where T: Integer,
+{
+    /// For a given n, iterate over all binomial coefficients binomial(n, k), for k=0...n.
+    ///
+    /// Note that this might overflow, depending on `T`. For the primitive
+    /// integer types, the following n are the largest ones for which there will
+    /// be no overflow:
+    ///
+    /// type | n
+    /// -----|---
+    /// u8   | 10
+    /// i8   |  9
+    /// u16  | 18
+    /// i16  | 17
+    /// u32  | 34
+    /// i32  | 33
+    /// u64  | 67
+    /// i64  | 66
+    ///
+    /// For larger n, `T` should be a bigint type.
+    pub fn new(n: T) -> IterBinomial<T> {
+        IterBinomial {
+            k: T::zero(), a: T::one(), n: n
+        }
+    }
+}
+
+impl<T> Iterator for IterBinomial<T>
+    where T: Integer + Clone
+{
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        if self.k > self.n {
+            return None;
+        }
+        self.a = if !self.k.is_zero() {
+            multiply_and_divide(
+                self.a.clone(),
+                self.n.clone() - self.k.clone() + T::one(),
+                self.k.clone()
+            )
+        } else {
+            T::one()
+        };
+        self.k = self.k.clone() + T::one();
+        Some(self.a.clone())
+    }
+}
+
 /// Calculate r * a / b, avoiding overflows and fractions.
+///
+/// Assumes that b divides r * a evenly.
 fn multiply_and_divide<T: Integer + Clone>(r: T, a: T, b: T) -> T {
     // See http://blog.plover.com/math/choose-2.html for the idea.
     let g = gcd(r.clone(), b.clone());
-    (r/g.clone() * a) / (b/g)
+    r/g.clone() * (a / (b/g))
 }
 
 /// Calculate the binomial coefficient.
+///
+/// Note that this might overflow, depending on `T`. For the primitive integer
+/// types, the following n are the largest ones possible such that there will
+/// be no overflow for any k:
+///
+/// type | n
+/// -----|---
+/// u8   | 10
+/// i8   |  9
+/// u16  | 18
+/// i16  | 17
+/// u32  | 34
+/// i32  | 33
+/// u64  | 67
+/// i64  | 66
+///
+/// For larger n, consider using a bigint type for `T`.
 pub fn binomial<T: Integer + Clone>(mut n: T, k: T) -> T {
     // See http://blog.plover.com/math/choose.html for the idea.
     if k > n {
@@ -737,6 +814,49 @@ fn test_lcm_overflow() {
     check!(u64, 0x8000_0000_0000_0000, 0x02, 0x8000_0000_0000_0000);
 }
 
+#[test]
+fn test_iter_binomial() {
+    macro_rules! check_simple {
+        ($t:ty) => { {
+            let n: $t = 3;
+            let c: Vec<_> = IterBinomial::new(n).collect();
+            let expected = vec![1, 3, 3, 1];
+            assert_eq!(c, expected);
+        } }
+    }
+
+    check_simple!(u8);
+    check_simple!(i8);
+    check_simple!(u16);
+    check_simple!(i16);
+    check_simple!(u32);
+    check_simple!(i32);
+    check_simple!(u64);
+    check_simple!(i64);
+
+    macro_rules! check_binomial {
+        ($t:ty, $n:expr) => { {
+            let n: $t = $n;
+            let c: Vec<_> = IterBinomial::new(n).collect();
+            let mut k: $t = 0;
+            for b in c {
+                assert_eq!(b, binomial(n, k));
+                k += 1;
+            }
+        } }
+    }
+
+    // Check the largest n for which there is no overflow.
+    check_binomial!(u8, 10);
+    check_binomial!(i8, 9);
+    check_binomial!(u16, 18);
+    check_binomial!(i16, 17);
+    check_binomial!(u32, 34);
+    check_binomial!(i32, 33);
+    check_binomial!(u64, 67);
+    check_binomial!(i64, 66);
+}
+
 #[test]
 fn test_binomial() {
     macro_rules! check {