Skip to content

Commit

Permalink
fix factorials
Browse files Browse the repository at this point in the history
  • Loading branch information
@ngtkana
ngtkana committed Jul 14, 2024
1 parent e45ebe1 commit e419e10
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 76 deletions.
166 changes: 93 additions & 73 deletions libs/fp/src/factorial.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,22 +76,26 @@ impl<const P: u64> Factorial<P> {
self.inv_fact[n]
}

/// The permutation $P(n, k)$
/// $[x^{n-k}]D^k x^n$.
/// # Requirements
/// - $k \le n \le \text{length}$
/// - $n \le \text{length}$
/// # Examples
/// ```
/// use fp::Factorial;
/// use fp::Fp;
/// const P: u64 = 998244353;
/// let fact = Factorial::<P>::new(10);
/// assert_eq!(fact.perm(8, 3).value(), 336);
/// assert_eq!(fact.falling(8, 3).value(), 336);
/// ```
pub fn perm(&self, n: usize, k: usize) -> Fp<P> {
self.fact[n] * self.inv_fact[n - k]
pub fn falling(&self, n: usize, k: usize) -> Fp<P> {
if n < k {
Fp::new(0)
} else {
self.fact[n] * self.inv_fact[n - k]
}
}

/// The binominal coefficient $n \choose k$
/// $[x^k](1 + x)^n$.
/// # Requirements
/// - $k \le n \le \text{length}$
/// # Examples
Expand All @@ -100,13 +104,18 @@ impl<const P: u64> Factorial<P> {
/// use fp::Fp;
/// const P: u64 = 998244353;
/// let fact = Factorial::<P>::new(10);
/// assert_eq!(fact.comb(8, 3).value(), 56);
/// assert_eq!(fact.binom(8, 3).value(), 56);
/// ```
pub fn comb(&self, n: usize, k: usize) -> Fp<P> {
self.fact[n] * self.inv_fact[n - k] * self.inv_fact[k]
pub fn binom(&self, n: usize, k: usize) -> Fp<P> {
assert!(n < self.fact.len());
if k > n {
Fp::new(0)
} else {
self.fact[n] * self.inv_fact[n - k] * self.inv_fact[k]
}
}

/// Another name for [`Factorial::comb`].
/// $[x^{a_1} \cdot x^{a_l}](x_{a_1} + \cdot x_{a_l})^n$.
/// # Requirements
/// - $k \le n \le \text{length}$
/// # Examples
Expand All @@ -115,13 +124,15 @@ impl<const P: u64> Factorial<P> {
/// use fp::Fp;
/// const P: u64 = 998244353;
/// let fact = Factorial::<P>::new(10);
/// assert_eq!(fact.binom(8, 3).value(), 56);
/// assert_eq!(fact.multinom(&[3, 5]).value(), 56);
/// ```
pub fn binom(&self, n: usize, k: usize) -> Fp<P> {
self.comb(n, k)
pub fn multinom(&self, a: &[usize]) -> Fp<P> {
let n = a.iter().sum::<usize>();
assert!(n < self.fact.len());
self.fact[n] * a.iter().map(|&x| self.inv_fact[x]).product::<Fp<P>>()
}

/// The binominal coefficient $n \choose k$, but zero if $k < 0$ or $k > n$
/// $[x^k](1 + x)^n$.
/// # Requirements
/// - $n \le \text{length}$
/// # Examples
Expand All @@ -130,18 +141,23 @@ impl<const P: u64> Factorial<P> {
/// use fp::Fp;
/// const P: u64 = 998244353;
/// let fact = Factorial::<P>::new(10);
/// assert_eq!(fact.comb_or_zero(8, 3).value(), 56);
/// assert_eq!(fact.comb_or_zero(8, 9).value(), 0);
/// assert_eq!(fact.binom_signed(8, 3).value(), 56);
/// assert_eq!(fact.binom_signed(8, 9).value(), 0);
/// ```
pub fn comb_or_zero(&self, n: usize, k: isize) -> Fp<P> {
if k < 0 || k as usize > n {
Fp::<P>::new(0)
pub fn binom_signed(&self, n: usize, k: isize) -> Fp<P> {
assert!(n < self.fact.len());
if k < 0 {
return Fp::new(0);
}
let k = k as usize;
if n < k {
Fp::new(0)
} else {
self.comb(n, k as usize)
self.fact[n] * self.inv_fact[n - k] * self.inv_fact[k]
}
}

/// The number of $k$-multicombinations of $n$ objects
/// $[x^k](1 - x)^{-n}$.
/// # Requirements
/// - $k \gt 0$ or $n \gt 0$
/// # Examples
Expand All @@ -150,11 +166,14 @@ impl<const P: u64> Factorial<P> {
/// use fp::Fp;
/// const P: u64 = 998244353;
/// let fact = Factorial::<P>::new(10);
/// assert_eq!(fact.comb_with_reputation(8, 3).value(), 120);
/// assert_eq!(fact.multiset_number(8, 3).value(), 120);
/// ```
pub fn comb_with_reputation(&self, n: usize, k: usize) -> Fp<P> {
assert!(n > 0 || k > 0);
self.comb(n + k - 1, k)
pub fn multiset_number(&self, n: usize, k: usize) -> Fp<P> {
if (n, k) == (0, 0) {
Fp::new(1)
} else {
self.binom(n + k - 1, k)
}
}
}
impl<const P: u64> Index<usize> for Factorial<P> {
Expand Down Expand Up @@ -190,69 +209,70 @@ mod tests {
}

#[test]
fn test_perm() {
fn test_falling() {
let fact = Factorial::<P>::new(10);
assert_eq!(fact.perm(0, 0).value(), 1);
assert_eq!(fact.perm(1, 0).value(), 1);
assert_eq!(fact.perm(1, 1).value(), 1);
assert_eq!(fact.perm(2, 0).value(), 1);
assert_eq!(fact.perm(2, 1).value(), 2);
assert_eq!(fact.perm(2, 2).value(), 2);
assert_eq!(fact.perm(3, 0).value(), 1);
assert_eq!(fact.perm(3, 1).value(), 3);
assert_eq!(fact.perm(3, 2).value(), 6);
assert_eq!(fact.perm(3, 3).value(), 6);
assert_eq!(fact.falling(0, 0).value(), 1);
assert_eq!(fact.falling(1, 0).value(), 1);
assert_eq!(fact.falling(1, 1).value(), 1);
assert_eq!(fact.falling(2, 0).value(), 1);
assert_eq!(fact.falling(2, 1).value(), 2);
assert_eq!(fact.falling(2, 2).value(), 2);
assert_eq!(fact.falling(3, 0).value(), 1);
assert_eq!(fact.falling(3, 1).value(), 3);
assert_eq!(fact.falling(3, 2).value(), 6);
assert_eq!(fact.falling(3, 3).value(), 6);
}

#[test]
fn test_comb() {
let fact = Factorial::<P>::new(10);
assert_eq!(fact.comb(0, 0).value(), 1);
assert_eq!(fact.comb(1, 0).value(), 1);
assert_eq!(fact.comb(1, 1).value(), 1);
assert_eq!(fact.comb(2, 0).value(), 1);
assert_eq!(fact.comb(2, 1).value(), 2);
assert_eq!(fact.comb(2, 2).value(), 1);
assert_eq!(fact.comb(3, 0).value(), 1);
assert_eq!(fact.comb(3, 1).value(), 3);
assert_eq!(fact.comb(3, 2).value(), 3);
assert_eq!(fact.comb(3, 3).value(), 1);
assert_eq!(fact.binom(0, 0).value(), 1);
assert_eq!(fact.binom(1, 0).value(), 1);
assert_eq!(fact.binom(1, 1).value(), 1);
assert_eq!(fact.binom(2, 0).value(), 1);
assert_eq!(fact.binom(2, 1).value(), 2);
assert_eq!(fact.binom(2, 2).value(), 1);
assert_eq!(fact.binom(3, 0).value(), 1);
assert_eq!(fact.binom(3, 1).value(), 3);
assert_eq!(fact.binom(3, 2).value(), 3);
assert_eq!(fact.binom(3, 3).value(), 1);
}

#[test]
fn test_comb_or_zero() {
let fact = Factorial::<P>::new(10);
assert_eq!(fact.comb_or_zero(0, -1).value(), 0);
assert_eq!(fact.comb_or_zero(0, 0).value(), 1);
assert_eq!(fact.comb_or_zero(0, 1).value(), 0);
assert_eq!(fact.comb_or_zero(1, -1).value(), 0);
assert_eq!(fact.comb_or_zero(1, 0).value(), 1);
assert_eq!(fact.comb_or_zero(1, 1).value(), 1);
assert_eq!(fact.comb_or_zero(1, 2).value(), 0);
assert_eq!(fact.comb_or_zero(2, -1).value(), 0);
assert_eq!(fact.comb_or_zero(2, 0).value(), 1);
assert_eq!(fact.comb_or_zero(2, 1).value(), 2);
assert_eq!(fact.comb_or_zero(2, 2).value(), 1);
assert_eq!(fact.comb_or_zero(2, 3).value(), 0);
assert_eq!(fact.comb_or_zero(3, -1).value(), 0);
assert_eq!(fact.comb_or_zero(3, 0).value(), 1);
assert_eq!(fact.comb_or_zero(3, 1).value(), 3);
assert_eq!(fact.comb_or_zero(3, 2).value(), 3);
assert_eq!(fact.comb_or_zero(3, 3).value(), 1);
assert_eq!(fact.comb_or_zero(3, 4).value(), 0);
assert_eq!(fact.binom_signed(0, -1).value(), 0);
assert_eq!(fact.binom_signed(0, 0).value(), 1);
assert_eq!(fact.binom_signed(0, 1).value(), 0);
assert_eq!(fact.binom_signed(1, -1).value(), 0);
assert_eq!(fact.binom_signed(1, 0).value(), 1);
assert_eq!(fact.binom_signed(1, 1).value(), 1);
assert_eq!(fact.binom_signed(1, 2).value(), 0);
assert_eq!(fact.binom_signed(2, -1).value(), 0);
assert_eq!(fact.binom_signed(2, 0).value(), 1);
assert_eq!(fact.binom_signed(2, 1).value(), 2);
assert_eq!(fact.binom_signed(2, 2).value(), 1);
assert_eq!(fact.binom_signed(2, 3).value(), 0);
assert_eq!(fact.binom_signed(3, -1).value(), 0);
assert_eq!(fact.binom_signed(3, 0).value(), 1);
assert_eq!(fact.binom_signed(3, 1).value(), 3);
assert_eq!(fact.binom_signed(3, 2).value(), 3);
assert_eq!(fact.binom_signed(3, 3).value(), 1);
assert_eq!(fact.binom_signed(3, 4).value(), 0);
}

#[test]
fn test_comb_with_reputation() {
let fact = Factorial::<P>::new(10);
assert_eq!(fact.comb_with_reputation(1, 0).value(), 1);
assert_eq!(fact.comb_with_reputation(1, 1).value(), 1);
assert_eq!(fact.comb_with_reputation(2, 0).value(), 1);
assert_eq!(fact.comb_with_reputation(2, 1).value(), 2);
assert_eq!(fact.comb_with_reputation(2, 2).value(), 3);
assert_eq!(fact.comb_with_reputation(3, 0).value(), 1);
assert_eq!(fact.comb_with_reputation(3, 1).value(), 3);
assert_eq!(fact.comb_with_reputation(3, 2).value(), 6);
assert_eq!(fact.comb_with_reputation(3, 3).value(), 10);
assert_eq!(fact.multiset_number(0, 0).value(), 1);
assert_eq!(fact.multiset_number(1, 0).value(), 1);
assert_eq!(fact.multiset_number(1, 1).value(), 1);
assert_eq!(fact.multiset_number(2, 0).value(), 1);
assert_eq!(fact.multiset_number(2, 1).value(), 2);
assert_eq!(fact.multiset_number(2, 2).value(), 3);
assert_eq!(fact.multiset_number(3, 0).value(), 1);
assert_eq!(fact.multiset_number(3, 1).value(), 3);
assert_eq!(fact.multiset_number(3, 2).value(), 6);
assert_eq!(fact.multiset_number(3, 3).value(), 10);
}
}
6 changes: 3 additions & 3 deletions libs/fp/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@
//! let f = Factorial::<998244353>::new(10);
//! assert_eq!(f.fact(5), fp!(120));
//! assert_eq!(f.inv_fact(5), fp!(120).inv());
//! assert_eq!(f.perm(5, 3), fp!(60));
//! assert_eq!(f.comb(5, 3), fp!(10));
//! assert_eq!(f.comb_with_reputation(5, 3), fp!(35));
//! assert_eq!(f.falling(5, 3), fp!(60));
//! assert_eq!(f.binom(5, 3), fp!(10));
//! assert_eq!(f.multiset_number(5, 3), fp!(35));
//! ```
//!
//! ## Convolution by Fast Fourier transform (FFT)
Expand Down

0 comments on commit e419e10

Please sign in to comment.