algonium 0.1.0

a collection of classic algorithms for programming contests
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use super::mint::{Mint, Module};

/// Useful struct to compute combinations
///
/// # Examples
/// ```
/// use algonium::math::{Comb, Mod107, Mint107};
/// let comb: Comb<Mod107> = Comb::new(100);
/// assert_eq!(Mint107::from(24), comb.fact(4));
/// assert_eq!(Mint107::from(1), comb.fact(4) * comb.factinv(4));
/// assert_eq!(Mint107::from(12), comb.perm(4, 2));
/// assert_eq!(Mint107::from(6), comb.comb(4, 2));
/// assert_eq!(Mint107::from(10), comb.multi_comb(4, 2));
/// ```
pub struct Comb<M: Module> {
    fact: Vec<Mint<M>>,
    factinv: Vec<Mint<M>>,
}

impl<M: Module> Comb<M> {
    /// Create a object that provides effiecint computation of combinations
    /// for input smaller than `n`.
    ///
    /// This requires `O(n)` time.
    pub fn new(n: usize) -> Comb<M> {
        let mut fact: Vec<Mint<M>> = vec![0.into(); n + 1];
        let mut factinv: Vec<Mint<M>> = vec![0.into(); n + 1];
        fact[0] = 1.into();
        for i in 0..n {
            fact[i + 1] = fact[i] * (i + 1);
        }
        factinv[n] = fact[n].inv();
        for i in (0..n).rev() {
            factinv[i] = factinv[i + 1] * (i + 1);
        }
        Comb {
            fact: fact,
            factinv: factinv,
        }
    }

    /// `n! = 1 * 2 * ... * n`
    ///
    /// `O(1)` if n is smaller than input in `new` method.
    pub fn fact(&self, n: u64) -> Mint<M> {
        if let Some(x) = self.fact.get(n as usize) {
            *x
        } else if n >= M::module() as u64 {
            Mint::from(0)
        } else {
            // Note that this is slow if `n` is large.
            // Precalculation is a possible solution but doesn't work for any module number.
            let mut res = 1.into();
            for a in 1..(n + 1) {
                res *= a;
            }
            res
        }
    }

    /// returns `y` such that `n! * y == 1`.
    ///
    /// `O(1)` if n is smaller than input in `new` method.
    pub fn factinv(&self, n: u64) -> Mint<M> {
        if let Some(x) = self.factinv.get(n as usize) {
            *x
        } else {
            self.fact(n).inv()
        }
    }

    /// `nPr = n! / (n - r)!`
    ///
    /// `O(1)` if n and r are smaller than input in `new` method.
    pub fn perm(&self, n: u64, r: u64) -> Mint<M> {
        if n >= r {
            self.fact(n) * self.factinv((n - r) as u64)
        } else {
            0.into()
        }
    }

    /// `nCr = n! / (n - r)! / r!`.
    ///
    /// `O(1)` if n and r are smaller than input in `new` method.
    pub fn comb(&self, n: u64, r: u64) -> Mint<M> {
        let m = M::module() as u64;
        if n >= m {
            self.comb(n % m, r % m) * self.comb(n / m, r / m) // Lucas' theorem
        } else if n >= r {
            self.fact(n) * self.factinv(n - r) * self.factinv(r)
        } else {
            Mint::from(0)
        }
    }

    /// `(n + k - 1)! / k!`.
    ///
    /// `O(1)` if n and r are smaller than input in `new` method.
    pub fn multi_comb(&self, n: u64, r: u64) -> Mint<M> {
        if r == 0 {
            Mint::from(1)
        } else {
            self.comb(n + r - 1, r)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple() {
        #[derive(Clone, Copy, Debug)]
        struct Mod;
        impl Module for Mod {
            fn module() -> u32 {
                1000000007
            }
        }
        let c = Comb::<Mod>::new(100);
        assert_eq!(Mint::from(336), c.perm(8, 3));
        assert_eq!(Mint::from(56), c.comb(8, 3));
        assert_eq!(Mint::from(10), c.multi_comb(3, 3));
    }

    #[test]
    fn test_fact() {
        #[derive(Clone, Copy, Debug)]
        struct Mod;
        impl Module for Mod {
            fn module() -> u32 {
                1000000007
            }
        }
        let c = Comb::<Mod>::new(100);
        let p = 8721234;
        let mut f = Mint::from(1);
        for i in 1..(p + 1) {
            f *= i;
        }
        assert_eq!(f, c.fact(p));
    }
}