Safe modular arithmetic and number theoretic transform

Author: EbTech

.travis.yml

@@ -1,6 +1,6 @@
 language: rust
 rust:
- - 1.31.1 # Version currently supported by Codeforces
+ # - 1.31.1 # Version currently supported by Codeforces
  - stable
  - beta
  - nightly

README.md

@@ -21,13 +21,15 @@ Most competition programmers rely on C++ for its fast execution time. However, i
 
 To my delight, I found that Rust eliminates entire classes of bugs, while reducing visual clutter to make the rest easier to spot. And, it's *fast*. There's a learning curve, to be sure. However, a proficient Rust programmer stands to gain a competitive advantage as well as a more pleasant experience!
 
+For help in getting started, you may check out [some of my past submissions](https://codeforces.com/contest/1168/submission/54859899).
+
 Other online judges that support Rust:
 - [Timus](http://acm.timus.ru/help.aspx?topic=rust)
 - [SPOJ](http://www.spoj.com/)
 
 ## Programming Language Advocacy
 
-My other goal is to appeal to developers who feel limited by mainstream, arguably outdated, programming languages. Perhaps we have a better way.
+My other goal is to appeal to developers who feel limited by mainstream, arguably outdated, programming languages. Perhaps we have a better option.
 
 Rather than try to persuade you with words, this repository aims to show by example. If you're new to Rust, see [Jim Blandy's *Why Rust?*](http://www.oreilly.com/programming/free/files/why-rust.pdf) for a brief introduction, or just [dive in!](https://doc.rust-lang.org/book/2018-edition)
 
@@ -37,6 +39,8 @@ Rather than try to persuade you with words, this repository aims to show by exam
 - [Network flows](src/graph/flow.rs): Dinic's blocking flow, Hopcroft-Karp bipartite matching, min cost max flow
 - [Connected components](src/graph/connectivity.rs): 2-edge-, 2-vertex- and strongly connected components, bridges, articulation points, topological sort, 2-SAT
 - [Associative range query](src/arq_tree.rs): known colloquially as *segtrees*
-- [Math](src/math/): Euclid's GCD algorithm, Bezout's identity, rational and complex numbers, fast Fourier transform
+- [GCD Math](src/math/mod.rs): canonical solution to Bezout's identity
+- [Arithmetic](src/math/num.rs): rational and complex numbers, safe modular arithmetic
+- [FFT](src/math/fft.rs): fast Fourier transform, number theoretic transform, convolution
 - [Scanner](src/scanner.rs): utility for reading input data
 - [String processing](src/string_proc.rs): Knuth-Morris-Pratt string matching, suffix arrays, Manacher's palindrome search

src/math/fft.rs

@@ -1,5 +1,6 @@
 //! The Fast Fourier Transform (FFT)
-use super::num::{Complex, PI};
+use super::num::{Complex, Field, PI};
+use std::ops::{Add, Div, Mul, Neg, Sub};
 
 // We can delete this struct once f64::reverse_bits() stabilizes.
 struct BitRevIterator {
@@ -29,25 +30,95 @@ impl Iterator for BitRevIterator {
  }
 }
 
-// Integer FFT notes: see problem 30-6 in CLRS for details, noting that
-// 15311432 and 469870224 are inverses and 2^23rd roots of 1 mod p=(119<<23)+1
-// 440564289 and 1713844692 are inverses and 2^27th roots of 1 mod p=(15<<27)+1
-// 125 and 2267742733 are inverses and 2^30th roots of 1 mod p=(3<<30)+1
+pub trait FFT: Sized + Copy {
+ type F: Sized
+ + Copy
+ + From<Self>
+ + Neg
+ + Add<Output = Self::F>
+ + Div<Output = Self::F>
+ + Mul<Output = Self::F>
+ + Sub<Output = Self::F>;
+
+ const ZERO: Self;
+
+ /// A primitive nth root of one raised to the powers 0, 1, 2, ..., n/2 - 1
+ fn get_roots(n: usize, inverse: bool) -> Vec<Self::F>;
+ /// 1 for forward transform, 1/n for inverse transform
+ fn get_factor(n: usize, inverse: bool) -> Self::F;
+ /// The inverse of Self::F::from()
+ fn extract(f: Self::F) -> Self;
+}
+
+impl FFT for f64 {
+ type F = Complex;
+
+ const ZERO: f64 = 0.0;
+
+ fn get_roots(n: usize, inverse: bool) -> Vec<Self::F> {
+ let step = if inverse { -2.0 } else { 2.0 } * PI / n as f64;
+ (0..n / 2)
+ .map(|i| Complex::from_polar(1.0, step * i as f64))
+ .collect()
+ }
+
+ fn get_factor(n: usize, inverse: bool) -> Self::F {
+ Self::F::from(if inverse { (n as f64).recip() } else { 1.0 })
+ }
+
+ fn extract(f: Self::F) -> f64 {
+ f.real
+ }
+}
+
+// NTT notes: see problem 30-6 in CLRS for details, keeping in mind that
+// 15311432 and 469870224 are inverses and 2^23rd roots of 1 mod (119<<23)+1
+// 440564289 and 1713844692 are inverses and 2^27th roots of 1 mod (15<<27)+1
+// 125 and 2267742733 are inverses and 2^30th roots of 1 mod (3<<30)+1
+impl FFT for u64 {
+ type F = Field;
+
+ const ZERO: u64 = 0;
+
+ fn get_roots(n: usize, inverse: bool) -> Vec<Self::F> {
+ assert!(n <= 1 << 23);
+ let mut prim_root = Self::F::from(15_311_432);
+ if inverse {
+ prim_root = prim_root.recip();
+ }
+ for _ in (0..).take_while(|&i| n < 1 << (23 - i)) {
+ prim_root = prim_root * prim_root;
+ }
+
+ let mut roots = Vec::with_capacity(n / 2);
+ let mut root = Self::F::from(1);
+ for _ in 0..roots.capacity() {
+ roots.push(root);
+ root = root * prim_root;
+ }
+ roots
+ }
+
+ fn get_factor(n: usize, inverse: bool) -> Self::F {
+ Self::F::from(if inverse { n as u64 } else { 1 }).recip()
+ }
+
+ fn extract(f: Self::F) -> u64 {
+ f.val
+ }
+}
 
 /// Computes the discrete fourier transform of v, whose length is a power of 2.
 /// Forward transform: polynomial coefficients -> evaluate at roots of unity
 /// Inverse transform: values at roots of unity -> interpolated coefficients
-pub fn fft(v: &[Complex], inverse: bool) -> Vec<Complex> {
+pub fn fft<T: FFT>(v: &[T::F], inverse: bool) -> Vec<T::F> {
  let n = v.len();
  assert!(n.is_power_of_two());
 
- let step = if inverse { -2.0 } else { 2.0 } * PI / n as f64;
- let factor = Complex::from(if inverse { n as f64 } else { 1.0 });
- let roots_of_unity = (0..n / 2)
- .map(|i| Complex::from_polar(1.0, step * i as f64))
- .collect::<Vec<_>>();
+ let factor = T::get_factor(n, inverse);
+ let roots_of_unity = T::get_roots(n, inverse);
  let mut dft = BitRevIterator::new(n)
- .map(|i| v[i] / factor)
+ .map(|i| v[i] * factor)
  .collect::<Vec<_>>();
 
  for m in (0..).map(|s| 1 << s).take_while(|&m| m < n) {
@@ -63,33 +134,36 @@ pub fn fft(v: &[Complex], inverse: bool) -> Vec<Complex> {
  dft
 }
 
-/// From a real vector, computes a DFT of size at least desired_len
-pub fn dft_from_reals(v: &[f64], desired_len: usize) -> Vec<Complex> {
+/// From a slice of reals (f64 or u64), computes DFT of size at least desired_len
+pub fn dft_from_reals<T: FFT>(v: &[T], desired_len: usize) -> Vec<T::F> {
  assert!(v.len() <= desired_len);
+ 
  let complex_v = v
  .iter()
  .cloned()
- .chain(std::iter::repeat(0.0))
+ .chain(std::iter::repeat(T::ZERO))
  .take(desired_len.next_power_of_two())
- .map(Complex::from)
+ .map(T::F::from)
  .collect::<Vec<_>>();
- fft(&complex_v, false)
+ fft::<T>(&complex_v, false)
 }
 
 /// The inverse of dft_from_reals()
-pub fn idft_to_reals(dft_v: &[Complex], desired_len: usize) -> Vec<f64> {
+pub fn idft_to_reals<T: FFT>(dft_v: &[T::F], desired_len: usize) -> Vec<T> {
  assert!(dft_v.len() >= desired_len);
- let complex_v = fft(dft_v, true);
+ 
+ let complex_v = fft::<T>(dft_v, true);
  complex_v
  .into_iter()
  .take(desired_len)
- .map(|c| c.real) // to get integers: c.real.round() as i64
+ .map(T::extract)
  .collect()
 }
 
-/// Given two polynomials (vectors) sum_i a[i]x^i and sum_i b[i]x^i,
-/// computes their product (convolution) c[k] = sum_(i+j=k) a[i]*b[j]
-pub fn convolution(a: &[f64], b: &[f64]) -> Vec<f64> {
+/// Given two polynomials (vectors) sum_i a[i] x^i and sum_i b[i] x^i,
+/// computes their product (convolution) c[k] = sum_(i+j=k) a[i]*b[j].
+/// Uses complex FFT if inputs are f64, or modular NTT if inputs are u64.
+pub fn convolution<T: FFT>(a: &[T], b: &[T]) -> Vec<T> {
  let len_c = a.len() + b.len() - 1;
  let dft_a = dft_from_reals(a, len_c).into_iter();
  let dft_b = dft_from_reals(b, len_c).into_iter();
@@ -102,10 +176,10 @@ mod test {
  use super::*;
 
  #[test]
- fn test_dft() {
+ fn test_complex_dft() {
  let v = vec![7.0, 1.0, 1.0];
  let dft_v = dft_from_reals(&v, v.len());
- let new_v = idft_to_reals(&dft_v, v.len());
+ let new_v: Vec<f64> = idft_to_reals(&dft_v, v.len());
 
  let six = Complex::from(6.0);
  let seven = Complex::from(7.0);
@@ -117,11 +191,40 @@ mod test {
  }
 
  #[test]
- fn test_convolution() {
+ fn test_modular_dft() {
+ let v = vec![7, 1, 1];
+ let dft_v = dft_from_reals(&v, v.len());
+ let new_v: Vec<u64> = idft_to_reals(&dft_v, v.len());
+
+ let seven = Field::from(7);
+ let one = Field::from(1);
+ let prim = Field::from(15_311_432).pow(1 << 21);
+ let prim2 = prim * prim;
+
+ let eval0 = seven + one + one;
+ let eval1 = seven + prim + prim2;
+ let eval2 = seven + prim2 + one;
+ let eval3 = seven + prim.recip() + prim2;
+
+ assert_eq!(dft_v, vec![eval0, eval1, eval2, eval3]);
+ assert_eq!(new_v, v);
+ }
+
+ #[test]
+ fn test_complex_convolution() {
  let x = vec![2.0, 3.0, 2.0];
  let y = vec![7.0, 2.0];
  let z = convolution(&x, &y);
 
  assert_eq!(z, vec![14.0, 25.0, 20.0, 4.0]);
  }
+
+ #[test]
+ fn test_modular_convolution() {
+ let x = vec![2, 3, 2];
+ let y = vec![7, 2];
+ let z = convolution(&x, &y);
+
+ assert_eq!(z, vec![14, 25, 20, 4]);
+ }
 }

src/math/mod.rs

@@ -2,30 +2,13 @@
 pub mod fft;
 pub mod num;
 
-/// Modular exponentiation by repeated squaring: returns base^exp % m.
-///
-/// # Panics
-///
-/// Panics if m == 0. May panic on overflow if m * m > 2^63.
-pub fn mod_pow(mut base: u64, mut exp: u64, m: u64) -> u64 {
- let mut result = 1 % m;
- while exp > 0 {
- if exp % 2 == 1 {
- result = (result * base) % m;
- }
- base = (base * base) % m;
- exp /= 2;
- }
- result
-}
-
 /// Finds (d, coef_a, coef_b) such that d = gcd(a, b) = a * coef_a + b * coef_b.
 pub fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) {
  if b == 0 {
  (a.abs(), a.signum(), 0)
  } else {
- let (d, coef_a, coef_b) = extended_gcd(b, a % b);
- (d, coef_b, coef_a - coef_b * (a / b))
+ let (d, coef_b, coef_a) = extended_gcd(b, a % b);
+ (d, coef_a, coef_b - coef_a * (a / b))
  }
 }
 
@@ -50,17 +33,6 @@ pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> {
 mod test {
  use super::*;
 
- #[test]
- fn test_mod_inverse() {
- let p = 1_000_000_007;
- let base = 31;
-
- let base_inv = mod_pow(base, p - 2, p);
- let identity = (base * base_inv) % p;
-
- assert_eq!(identity, 1);
- }
-
  #[test]
  fn test_egcd() {
  let (a, b) = (14, 35);