add binom (#20)

Author: cameroncuster

examples/binomial_coefficient_prime_mod.rs

@@ -0,0 +1,19 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/binomial_coefficient_prime_mod
+
+use proconio::input;
+use programming_team_code_rust::binom::Binom;
+
+fn main() {
+    input! {
+        t: usize,
+        m: u64,
+    }
+    let mut binom = Binom::new(m);
+    for _ in 0..t {
+        input! {
+            n: usize,
+            k: usize,
+        }
+        println!("{}", binom.comb(n, k));
+    }
+}

src/binom.rs

@@ -0,0 +1,34 @@
+// n choose k for a prime modulo
+pub struct Binom {
+    inv: Vec<u64>,
+    fact: Vec<u64>,
+    inv_fact: Vec<u64>,
+    modulo: u64,
+}
+impl Binom {
+    pub fn new(modulo: u64) -> Binom {
+        Binom {
+            inv: vec![1; 2],
+            fact: vec![1; 2],
+            inv_fact: vec![1; 2],
+            modulo,
+        }
+    }
+    pub fn comb(&mut self, n: usize, k: usize) -> u64 {
+        assert!((n as u64) < self.modulo);
+        if n < k {
+            return 0;
+        }
+        while self.inv.len() <= n {
+            let i = self.inv.len();
+            self.inv.push(
+                self.modulo
+                    - (self.modulo / i as u64) * self.inv[self.modulo as usize % i] % self.modulo,
+            );
+            self.fact.push(self.fact[i - 1] * i as u64 % self.modulo);
+            self.inv_fact
+                .push(self.inv_fact[i - 1] * self.inv[i] % self.modulo);
+        }
+        self.fact[n] * self.inv_fact[k] % self.modulo * self.inv_fact[n - k] % self.modulo
+    }
+}

src/lib.rs

@@ -1,3 +1,4 @@
+pub mod binom;
 pub mod dijk;
 pub mod dsu;
 pub mod fenwick;