contest_algorithms/math/
mod.rs

1//! Number-theoretic utilities for contest problems.
2pub mod fft;
3pub mod num;
4
5/// Finds (d, coef_a, coef_b) such that d = gcd(a, b) = a * coef_a + b * coef_b.
6pub fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) {
7    if b == 0 {
8        (a.abs(), a.signum(), 0)
9    } else {
10        let (d, coef_b, coef_a) = extended_gcd(b, a % b);
11        (d, coef_a, coef_b - coef_a * (a / b))
12    }
13}
14
15/// Assuming a != 0, finds smallest coef_b >= 0 such that a * coef_a + b * coef_b = c.
16///
17/// # Panics
18///
19/// Panics if a == 0.
20pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> {
21    let (d, _, coef_b_init) = extended_gcd(a, b);
22    if c % d == 0 {
23        let a_d = (a / d).abs();
24        let coef_b = (coef_b_init * (c / d) % a_d + a_d) % a_d;
25        let coef_a = (c - b * coef_b) / a;
26        Some((d, coef_a, coef_b))
27    } else {
28        None
29    }
30}
31
32// TODO: deduplicate modular arithmetic code with num::Field
33fn pos_mod(n: i64, m: i64) -> i64 {
34    if n < 0 {
35        n + m
36    } else {
37        n
38    }
39}
40fn mod_mul(a: i64, b: i64, m: i64) -> i64 {
41    pos_mod((a as i128 * b as i128 % m as i128) as i64, m)
42}
43fn mod_exp(mut base: i64, mut exp: u64, m: i64) -> i64 {
44    assert!(m >= 1);
45    let mut ans = 1 % m;
46    base %= m;
47    while exp > 0 {
48        if exp % 2 == 1 {
49            ans = mod_mul(ans, base, m);
50        }
51        base = mod_mul(base, base, m);
52        exp /= 2;
53    }
54    pos_mod(ans, m)
55}
56
57fn is_strong_probable_prime(n: i64, exp: u64, r: i64, a: i64) -> bool {
58    let mut x = mod_exp(a, exp, n);
59    if x == 1 || x == n - 1 {
60        return true;
61    }
62    for _ in 1..r {
63        x = mod_mul(x, x, n);
64        if x == n - 1 {
65            return true;
66        }
67    }
68    false
69}
70
71/// Assuming x >= 0, returns whether x is prime
72pub fn is_prime(n: i64) -> bool {
73    const BASES: [i64; 12] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37];
74    assert!(n >= 0);
75    match n {
76        0 | 1 => false,
77        2 | 3 => true,
78        _ if n % 2 == 0 => false,
79        _ => {
80            let r = (n - 1).trailing_zeros() as i64;
81            let exp = (n - 1) as u64 >> r;
82            BASES
83                .iter()
84                .all(|&base| base > n - 2 || is_strong_probable_prime(n, exp, r, base))
85        }
86    }
87}
88
89fn pollard_rho(n: i64) -> i64 {
90    for a in 1..n {
91        let f = |x| pos_mod(mod_mul(x, x, n) + a, n);
92        let mut x = 2;
93        let mut y = 2;
94        loop {
95            x = f(x);
96            y = f(f(y));
97            let div = num::fast_gcd(x - y, n);
98            if div == n {
99                break;
100            } else if div > 1 {
101                return div;
102            }
103        }
104    }
105    panic!("No divisor found!");
106}
107
108/// Assuming x >= 1, finds the prime factorization of n
109/// TODO: pollard_rho needs randomization to ensure correctness in contest settings!
110pub fn factorize(n: i64) -> Vec<i64> {
111    assert!(n >= 1);
112    let r = n.trailing_zeros() as usize;
113    let mut factors = vec![2; r];
114    let mut stack = match n >> r {
115        1 => vec![],
116        x => vec![x],
117    };
118    while let Some(top) = stack.pop() {
119        if is_prime(top) {
120            factors.push(top);
121        } else {
122            let div = pollard_rho(top);
123            stack.push(div);
124            stack.push(top / div);
125        }
126    }
127    factors.sort_unstable();
128    factors
129}
130
131#[cfg(test)]
132mod test {
133    use super::*;
134
135    #[test]
136    fn test_egcd() {
137        let (a, b) = (14, 35);
138
139        let (d, x, y) = extended_gcd(a, b);
140        assert_eq!(d, 7);
141        assert_eq!(a * x + b * y, d);
142
143        assert_eq!(canon_egcd(a, b, d), Some((d, -2, 1)));
144        assert_eq!(canon_egcd(b, a, d), Some((d, -1, 3)));
145    }
146
147    #[test]
148    fn test_modexp() {
149        let m = 1_000_000_007;
150        assert_eq!(mod_exp(0, 0, m), 1);
151        assert_eq!(mod_exp(0, 1, m), 0);
152        assert_eq!(mod_exp(0, 10, m), 0);
153        assert_eq!(mod_exp(123, 456, m), 565291922);
154    }
155
156    #[test]
157    fn test_miller() {
158        assert_eq!(is_prime(2), true);
159        assert_eq!(is_prime(4), false);
160        assert_eq!(is_prime(6), false);
161        assert_eq!(is_prime(8), false);
162        assert_eq!(is_prime(269), true);
163        assert_eq!(is_prime(1000), false);
164        assert_eq!(is_prime(1_000_000_007), true);
165        assert_eq!(is_prime((1 << 61) - 1), true);
166        assert_eq!(is_prime(7156857700403137441), false);
167    }
168
169    #[test]
170    fn test_pollard() {
171        assert_eq!(factorize(1), vec![]);
172        assert_eq!(factorize(2), vec![2]);
173        assert_eq!(factorize(4), vec![2, 2]);
174        assert_eq!(factorize(12), vec![2, 2, 3]);
175        assert_eq!(
176            factorize(7156857700403137441),
177            vec![11, 13, 17, 19, 29, 37, 41, 43, 61, 97, 109, 127]
178        );
179    }
180}
contest_algorithms/math/mod.rs

contest_algorithms/math/
mod.rs
