sim_lib_discrete_comb/
count.rs

1//! Exact combinatorial counting functions over `num_bigint::BigUint`.
2
3use num_bigint::BigUint;
4
5use crate::error::CombError;
6
7/// Largest `n` accepted by [`factorial_checked`] before [`CombError::LimitExceeded`].
8///
9/// `factorial` folds an unbounded range of growing big-integer products; an
10/// uncapped value parsed from untrusted text would hang or exhaust memory.
11pub const MAX_FACTORIAL_INPUT: u64 = 50_000;
12
13/// Largest `n` accepted by [`integer_partition_count_checked`] before
14/// [`CombError::LimitExceeded`].
15///
16/// `integer_partition_count` allocates `vec![..; n + 1]` and runs an `O(n^2)`
17/// big-integer recurrence, so an uncapped value is an out-of-memory / hang risk.
18pub const MAX_PARTITION_INPUT: u64 = 10_000;
19
20/// Largest effective `k` (`min(k, n - k)`) accepted by [`binomial_checked`]
21/// before [`CombError::LimitExceeded`].
22///
23/// `binomial` folds a `k`-length big-integer product loop; an uncapped `k`
24/// parsed from untrusted text would hang or exhaust memory. Matched to
25/// [`MAX_FACTORIAL_INPUT`], the sibling product-fold ceiling.
26pub const MAX_BINOMIAL_INPUT: u64 = MAX_FACTORIAL_INPUT;
27
28fn big(n: u64) -> BigUint {
29    BigUint::from(n)
30}
31
32/// `n!` (with `factorial(0) == 1`).
33pub fn factorial(n: u64) -> BigUint {
34    (1..=n).fold(BigUint::from(1u32), |acc, i| acc * big(i))
35}
36
37/// `n!` with an explicit input ceiling ([`MAX_FACTORIAL_INPUT`]).
38///
39/// Returns [`CombError::LimitExceeded`] for `n` beyond the cap instead of
40/// folding an unbounded range from untrusted input.
41pub fn factorial_checked(n: u64) -> Result<BigUint, CombError> {
42    if n > MAX_FACTORIAL_INPUT {
43        return Err(CombError::LimitExceeded(format!(
44            "factorial input {n} exceeds maximum {MAX_FACTORIAL_INPUT}"
45        )));
46    }
47    Ok(factorial(n))
48}
49
50/// The falling factorial `n * (n-1) * ... * (n-k+1)` (`0` when `k > n`).
51pub fn falling_factorial(n: u64, k: u64) -> BigUint {
52    if k > n {
53        return BigUint::from(0u32);
54    }
55    (0..k).fold(BigUint::from(1u32), |acc, i| acc * big(n - i))
56}
57
58/// The number of `k`-permutations of `n`, `nPk` (`0` when `k > n`).
59pub fn permutation_count(n: u64, k: u64) -> BigUint {
60    falling_factorial(n, k)
61}
62
63/// `n choose k` via the multiplicative formula (`0` when `k > n`).
64///
65/// # Examples
66///
67/// ```
68/// use num_bigint::BigUint;
69/// use sim_lib_discrete_comb::binomial;
70///
71/// assert_eq!(binomial(5, 2), BigUint::from(10u32));
72/// assert_eq!(binomial(5, 0), BigUint::from(1u32)); // empty choice
73/// assert_eq!(binomial(2, 5), BigUint::from(0u32)); // k > n
74/// ```
75pub fn binomial(n: u64, k: u64) -> BigUint {
76    if k > n {
77        return BigUint::from(0u32);
78    }
79    let k = k.min(n - k);
80    let mut result = BigUint::from(1u32);
81    for i in 1..=k {
82        // result is always divisible by i at this step, so division is exact.
83        result = result * big(n - k + i) / big(i);
84    }
85    result
86}
87
88/// `n choose k` with an explicit ceiling on the effective loop length
89/// (`min(k, n - k)`).
90///
91/// Returns [`CombError::LimitExceeded`] when `min(k, n - k)` exceeds `max_iter`
92/// instead of folding a `k`-length big-integer product from untrusted input.
93pub fn binomial_checked(n: u64, k: u64, max_iter: u64) -> Result<BigUint, CombError> {
94    let effective = k.min(n.saturating_sub(k));
95    if effective > max_iter {
96        return Err(CombError::LimitExceeded(format!(
97            "binomial k={effective} exceeds maximum {max_iter}"
98        )));
99    }
100    Ok(binomial(n, k))
101}
102
103/// The multinomial coefficient `(sum parts)! / prod(part!)`.
104pub fn multinomial(parts: &[u64]) -> BigUint {
105    let total: u64 = parts.iter().sum();
106    let mut denom = BigUint::from(1u32);
107    for &p in parts {
108        denom *= factorial(p);
109    }
110    factorial(total) / denom
111}
112
113/// Stirling numbers of the second kind `S(n, k)`: partitions of an `n`-set into
114/// `k` non-empty unlabeled blocks.
115pub fn stirling2(n: u64, k: u64) -> BigUint {
116    let (n, k) = (n as usize, k as usize);
117    let mut dp = vec![BigUint::from(0u32); k + 1];
118    dp[0] = BigUint::from(1u32); // S(0,0) = 1
119    for _ in 1..=n {
120        let mut next = vec![BigUint::from(0u32); k + 1];
121        for j in 1..=k {
122            next[j] = big(j as u64) * &dp[j] + &dp[j - 1];
123        }
124        dp = next;
125    }
126    if k < dp.len() {
127        dp[k].clone()
128    } else {
129        BigUint::from(0u32)
130    }
131}
132
133/// The Bell number `B(n) = sum_k S(n, k)`: total partitions of an `n`-set.
134pub fn bell_number(n: u64) -> BigUint {
135    (0..=n).map(|k| stirling2(n, k)).sum()
136}
137
138/// The partition count `p(n)`: ways to write `n` as a sum of positive integers,
139/// order ignored.
140pub fn integer_partition_count(n: u64) -> BigUint {
141    let n = n as usize;
142    let mut dp = vec![BigUint::from(0u32); n + 1];
143    dp[0] = BigUint::from(1u32);
144    for part in 1..=n {
145        for j in part..=n {
146            dp[j] = &dp[j] + &dp[j - part].clone();
147        }
148    }
149    dp[n].clone()
150}
151
152/// `p(n)` with an explicit input ceiling ([`MAX_PARTITION_INPUT`]).
153///
154/// Returns [`CombError::LimitExceeded`] for `n` beyond the cap, so an untrusted
155/// `n` cannot drive an unbounded allocation or `O(n^2)` recurrence.
156pub fn integer_partition_count_checked(n: u64) -> Result<BigUint, CombError> {
157    if n > MAX_PARTITION_INPUT {
158        return Err(CombError::LimitExceeded(format!(
159            "partition-count input {n} exceeds maximum {MAX_PARTITION_INPUT}"
160        )));
161    }
162    Ok(integer_partition_count(n))
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    fn b(n: u64) -> BigUint {
170        BigUint::from(n)
171    }
172
173    #[test]
174    fn factorial_basics() {
175        assert_eq!(factorial(0), b(1));
176        assert_eq!(factorial(5), b(120));
177    }
178
179    #[test]
180    fn binomial_and_permutations() {
181        assert_eq!(binomial(5, 2), b(10));
182        assert_eq!(binomial(10, 5), b(252));
183        assert_eq!(binomial(5, 7), b(0));
184        assert_eq!(binomial(6, 0), b(1));
185        assert_eq!(permutation_count(5, 3), b(60));
186        assert_eq!(falling_factorial(5, 2), b(20));
187    }
188
189    #[test]
190    fn multinomial_value() {
191        // 4! / (2! 1! 1!) = 12
192        assert_eq!(multinomial(&[2, 1, 1]), b(12));
193    }
194
195    #[test]
196    fn stirling_bell_and_partitions() {
197        assert_eq!(stirling2(4, 2), b(7));
198        assert_eq!(bell_number(4), b(15));
199        assert_eq!(integer_partition_count(5), b(7));
200    }
201
202    #[test]
203    fn checked_counts_accept_small_inputs() {
204        assert_eq!(factorial_checked(5).unwrap(), b(120));
205        assert_eq!(integer_partition_count_checked(5).unwrap(), b(7));
206        assert_eq!(binomial_checked(5, 2, MAX_BINOMIAL_INPUT).unwrap(), b(10));
207        // k > n stays a cheap zero, never rejected on the huge n - k underflow.
208        assert_eq!(binomial_checked(2, 5, MAX_BINOMIAL_INPUT).unwrap(), b(0));
209    }
210
211    #[test]
212    fn checked_counts_reject_huge_inputs() {
213        assert!(matches!(
214            factorial_checked(MAX_FACTORIAL_INPUT + 1),
215            Err(CombError::LimitExceeded(_))
216        ));
217        assert!(matches!(
218            integer_partition_count_checked(u64::MAX),
219            Err(CombError::LimitExceeded(_))
220        ));
221        // min(k, n - k) is the loop length, so a huge symmetric k is rejected.
222        assert!(matches!(
223            binomial_checked(1_000_000_000_000, 500_000_000_000, MAX_BINOMIAL_INPUT),
224            Err(CombError::LimitExceeded(_))
225        ));
226    }
227}
sim_lib_discrete_comb/count.rs

sim_lib_discrete_comb/
count.rs
