competitive/math/
prime_factors.rs

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

competitive/math/
prime_factors.rs
