algorithm_rs/
math.rs

1pub fn com(n: i64, mut r: i64) -> i64 {
2    if r > n {
3        return 0;
4    }
5
6    if r * 2 > n {
7        r = n - r;
8    }
9
10    if r == 0 {
11        return 1;
12    }
13
14    let mut res = 1;
15    for i in 1..(r + 1) {
16        res *= n - i + 1;
17        res /= i;
18    }
19
20    res
21}
22
23#[cfg(test)]
24mod test_com {
25    use crate::math::com;
26
27    #[test]
28    fn it_works() {
29        assert_eq!(com(6, 0), 1);
30        assert_eq!(com(6, 1), 6);
31        assert_eq!(com(6, 2), 15);
32        assert_eq!(com(6, 3), 20);
33        assert_eq!(com(6, 4), 15);
34        assert_eq!(com(6, 5), 6);
35        assert_eq!(com(6, 6), 1);
36    }
37}
38
39pub fn eratosthenes_sieve(n: usize) -> Vec<bool> {
40    let mut s = vec![true; n];
41    s[0] = false;
42    s[1] = false;
43
44    for i in 2..n {
45        if i * i > n {
46            break;
47        }
48        if s[i] {
49            for k in 2..(n + i - 1) / i {
50                s[k * i] = false
51            }
52        }
53    }
54
55    s
56}
57
58#[cfg(test)]
59mod test_eratosthenes_sieve {
60    use crate::math::eratosthenes_sieve;
61
62    #[test]
63    fn it_works() {
64        let primes = eratosthenes_sieve(10);
65
66        assert!(!primes[1]);
67        assert!(primes[2]);
68        assert!(primes[3]);
69        assert!(!primes[4]);
70        assert!(primes[5]);
71        assert!(!primes[6]);
72        assert!(primes[7]);
73        assert!(!primes[8]);
74        assert!(!primes[9]);
75    }
76}
77
78pub fn gcd(a: i64, b: i64) -> i64 {
79    if b == 0 {
80        a
81    } else {
82        gcd(b, a % b)
83    }
84}
85
86#[cfg(test)]
87mod test_gcd {
88    use crate::math::gcd;
89
90    #[test]
91    fn it_works() {
92        assert_eq!(gcd(15, 5), 5);
93        assert_eq!(gcd(5, 15), 5);
94        assert_eq!(gcd(198, 26), 2);
95        assert_eq!(gcd(26, 198), 2);
96    }
97}
98
99pub fn modcom(n: usize, k: usize, div: usize) -> usize {
100    let mut fact = Vec::<usize>::new();
101    let mut inv = Vec::<usize>::new();
102    let mut finv = Vec::<usize>::new();
103
104    let upper = n + 1;
105
106    fact.resize(upper, 0);
107    inv.resize(upper, 0);
108    finv.resize(upper, 0);
109
110    fact[0] = 1;
111    fact[1] = 1;
112
113    finv[0] = 1;
114    finv[1] = 1;
115
116    inv[1] = 1;
117
118    for i in 2..upper {
119        fact[i] = fact[i - 1] * i % div;
120        inv[i] = div - inv[div % i] * (div / i) % div;
121        finv[i] = finv[i - 1] * inv[i] % div;
122    }
123
124    fact[n] * (finv[k] * finv[n - k] % div) % div
125}
126
127#[cfg(test)]
128mod test_modcom {
129    use crate::math::modcom;
130
131    #[test]
132    fn it_works() {
133        assert_eq!(modcom(6, 0, 7), 1);
134        assert_eq!(modcom(6, 1, 7), 6);
135        assert_eq!(modcom(6, 2, 7), 1);
136        assert_eq!(modcom(6, 3, 7), 6);
137        assert_eq!(modcom(6, 4, 7), 1);
138        assert_eq!(modcom(6, 5, 7), 6);
139        assert_eq!(modcom(6, 6, 7), 1);
140    }
141}
142
143pub fn modpow(mut a: usize, mut n: usize, m: usize) -> usize {
144    let mut ans = 1;
145
146    while n > 0 {
147        if n & 1 == 1 {
148            ans = ans * a % m;
149        }
150
151        a = a * a % m;
152        n >>= 1;
153    }
154
155    ans
156}
157
158#[cfg(test)]
159mod test_modpow {
160    use crate::math::modpow;
161
162    #[test]
163    fn it_works() {
164        assert_eq!(modpow(2, 10, 1_000_000_007), 1024);
165        assert_eq!(modpow(2, 3, 1_000_000_007), 8);
166        assert_eq!(modpow(5, 8, 1_000_000_007), 390625);
167        assert_eq!(modpow(2, 2, 3), 1);
168        assert_eq!(modpow(2, 3, 3), 2);
169    }
170}
171
172pub fn phi(mut n: usize) -> usize {
173    let mut res = n;
174
175    let mut i = 2;
176    while i * i <= n {
177        if n % i == 0 {
178            res = res / i * (i - 1);
179            while n % i == 0 {
180                n /= i;
181            }
182        }
183        i += 1;
184    }
185
186    if n != 1 {
187        res = res / n * (n - 1);
188    }
189
190    res
191}
192
193pub fn modinv(a: usize, p: usize) -> usize {
194    let m = phi(p);
195    modpow(a, m - 1, p)
196}
197
198#[cfg(test)]
199mod test_modinv {
200    use crate::math::modinv;
201    use crate::math::modpow;
202
203    fn inv_euler(a: usize, p: usize) -> usize {
204        modpow(a, p - 2, p)
205    }
206
207    fn inv_simple(a: usize, p: usize) -> usize {
208        let mut ret = 1;
209        loop {
210            if (a * ret) % p == 1 {
211                return ret;
212            }
213            ret += 1;
214        }
215    }
216
217    #[test]
218    fn it_works_prime() {
219        // 10^-1 mod 7
220        assert_eq!(modinv(10, 7), inv_euler(10, 7));
221    }
222
223    #[test]
224    fn it_works_composite() {
225        // 10^-1 mod 2019
226        assert_eq!(modinv(10, 2019), inv_simple(10, 2019));
227    }
228}
229
230#[derive(Debug, Clone)]
231pub struct SequentialPrimeFactorization {
232    smallest_prime_factors: Vec<usize>,
233}
234
235impl SequentialPrimeFactorization {
236    pub fn new(n: usize) -> Self {
237        let mut smallest_prime_factors: Vec<usize> = (0..=(n + 1)).collect();
238        let mut i = 2;
239
240        while i * i <= n {
241            if smallest_prime_factors[i] == i {
242                let mut j = i;
243                while j * i <= n {
244                    smallest_prime_factors[j * i] = i;
245                    j += 1;
246                }
247            }
248
249            i += 1;
250        }
251
252        Self { smallest_prime_factors }
253    }
254
255    pub fn factorize(&self, mut x: usize) -> Vec<usize> {
256        let mut ret: Vec<usize> = vec![];
257        while x != 1 {
258            ret.push(self.smallest_prime_factors[x]);
259            x /= self.smallest_prime_factors[x];
260        }
261
262        ret.sort_unstable();
263        ret
264    }
265}
266
267pub fn prime_factorize(mut n: usize) -> Vec<(usize, usize)> {
268    let mut ret = vec![];
269    let mut p = 2;
270
271    while p * p <= n {
272        if n % p == 0 {
273            let mut k = 0;
274            while n % p == 0 {
275                k += 1;
276                n /= p;
277            }
278            ret.push((p, k));
279        }
280
281        p += 1;
282    }
283
284    if n != 1 {
285        ret.push((n, 1));
286    }
287    ret
288}
289
290#[cfg(test)]
291mod test_prime_factorization {
292    use crate::math::prime_factorize;
293    use crate::math::SequentialPrimeFactorization;
294
295    #[test]
296    fn it_works_sequential_prime_factorization() {
297        let prime_factorizer = SequentialPrimeFactorization::new(100);
298
299        assert_eq!(prime_factorizer.factorize(1), vec![]);
300        assert_eq!(prime_factorizer.factorize(2), vec![2]);
301        assert_eq!(prime_factorizer.factorize(4), vec![2, 2]);
302        assert_eq!(prime_factorizer.factorize(7), vec![7]);
303        assert_eq!(prime_factorizer.factorize(30), vec![2, 3, 5]);
304        assert_eq!(prime_factorizer.factorize(23), vec![23]);
305    }
306
307    #[test]
308    fn it_works_prime_factorization() {
309        assert_eq!(prime_factorize(1), vec![]);
310        assert_eq!(prime_factorize(2), vec![(2, 1)]);
311        assert_eq!(prime_factorize(4), vec![(2, 2)]);
312        assert_eq!(prime_factorize(7), vec![(7, 1)]);
313        assert_eq!(prime_factorize(30), vec![(2, 1), (3, 1), (5, 1)]);
314        assert_eq!(prime_factorize(23), vec![(23, 1)]);
315        assert_eq!(prime_factorize(512), vec![(2, 9)]);
316    }
317}