ソースを参照

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 {