competitive/math/
prime_factors.rs

1use super::{BarrettReduction, Xorshift, gcd, miller_rabin_with_br};
2
3struct MontgomeryReduction64 {
4    modulus: u64,
5    inverse: u64,
6}
7
8impl MontgomeryReduction64 {
9    fn new(modulus: u64) -> Self {
10        let mut inverse = modulus;
11        for _ in 0..6 {
12            inverse = inverse.wrapping_mul(2u64.wrapping_sub(modulus.wrapping_mul(inverse)));
13        }
14        Self { modulus, inverse }
15    }
16
17    fn sub(&self, lhs: u64, rhs: u64) -> u64 {
18        let (value, borrow) = lhs.overflowing_sub(rhs);
19        value.wrapping_add((borrow as u64).wrapping_neg() & self.modulus)
20    }
21
22    fn mul(&self, lhs: u64, rhs: u64) -> u64 {
23        let product = lhs as u128 * rhs as u128;
24        let (value, borrow) = ((product >> 64) as u64).overflowing_sub(
25            (((product as u64).wrapping_mul(self.inverse) as u128 * self.modulus as u128) >> 64)
26                as u64,
27        );
28        value.wrapping_add((borrow as u64).wrapping_neg() & self.modulus)
29    }
30}
31
32fn find_factor(n: u64) -> Option<u64> {
33    let br = BarrettReduction::<u128>::new(n as u128);
34    if miller_rabin_with_br(n, &br) {
35        return None;
36    }
37    let mr = MontgomeryReduction64::new(n);
38    let mut rng = Xorshift::default();
39    let (mut y0, mut c) = (0, n - 1);
40    loop {
41        let (mut x, mut y, mut ys, mut g, mut q, mut r, mut k) = (0, y0, 0, 1, 1, 1, 0);
42        while g == 1 && r <= 1 << 20 {
43            x = y;
44            while k < r && g == 1 {
45                ys = y;
46                for _ in 0..1024.min(r - k) {
47                    y = mr.sub(mr.mul(y, y), c);
48                    q = mr.mul(q, mr.sub(x, y));
49                }
50                g = gcd(q, n);
51                k += 1024;
52            }
53            k = r;
54            r <<= 1;
55        }
56        if g == n {
57            g = 1;
58            y = ys;
59            while g == 1 {
60                y = mr.sub(mr.mul(y, y), c);
61                g = gcd(mr.sub(x, y), n);
62            }
63        }
64        if g != 1 && g != n {
65            return Some(g);
66        }
67        y0 = ((rng.rand64() as u128 * (n - 2) as u128) >> 64) as u64 + 2;
68        c = ((rng.rand64() as u128 * (n - 1) as u128) >> 64) as u64 + 1;
69    }
70}
71
72pub fn prime_factors_flatten(mut n: u64) -> Vec<u64> {
73    if n == 0 {
74        return vec![];
75    }
76    let k = n.trailing_zeros();
77    let mut res = vec![2; k as usize];
78    n >>= k;
79    while n.is_multiple_of(3) {
80        res.push(3);
81        n /= 3;
82    }
83    if n != 1 {
84        let mut c = vec![n];
85        while let Some(n) = c.pop() {
86            if let Some(m) = find_factor(n) {
87                c.push(m);
88                c.push(n / m);
89            } else {
90                res.push(n);
91            }
92        }
93    }
94    res.sort_unstable();
95    res
96}
97
98pub fn prime_factors(n: u64) -> Vec<(u64, u32)> {
99    let mut res = Vec::new();
100    for a in prime_factors_flatten(n) {
101        if let Some((p, len)) = res.last_mut()
102            && p == &a
103        {
104            *len += 1;
105            continue;
106        }
107        res.push((a, 1));
108    }
109    res
110}
111
112pub fn divisors(n: u64) -> Vec<u64> {
113    let mut d = vec![1u64];
114    for (p, c) in prime_factors(n) {
115        let k = d.len();
116        let mut acc = 1;
117        for _ in 0..c {
118            acc *= p;
119            for i in 0..k {
120                d.push(d[i] * acc);
121            }
122        }
123    }
124    d.sort_unstable();
125    d
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use crate::tools::Xorshift;
132
133    pub fn naive_divisors(n: u64) -> Vec<u64> {
134        let mut res = vec![];
135        for i in 1..(n as f32).sqrt() as u64 + 1 {
136            if n.is_multiple_of(i) {
137                res.push(i);
138                if i * i != n {
139                    res.push(n / i);
140                }
141            }
142        }
143        res.sort_unstable();
144        res
145    }
146
147    #[test]
148    fn test_prime_factors_rho() {
149        use crate::{math::miller_rabin, tools::Xorshift};
150        const Q: usize = 2_000;
151        let mut rng = Xorshift::default();
152        for _ in 0..Q {
153            let x = rng.rand64();
154            let factors = prime_factors_flatten(x);
155            assert!(factors.iter().all(|&p| miller_rabin(p)));
156            let p = factors.into_iter().product::<u64>();
157            assert_eq!(x, p);
158        }
159    }
160
161    #[test]
162    fn test_divisors() {
163        let mut rng = Xorshift::default();
164        for n in (1..1000).chain(rng.random_iter(1..=20000000).take(100)) {
165            assert_eq!(divisors(n), naive_divisors(n));
166        }
167    }
168}
competitive/math/prime_factors.rs

competitive/math/
prime_factors.rs
