competitive/math/
prime_table.rs

1use std::iter::once;
2
3#[derive(Clone, Debug)]
4pub struct PrimeTable {
5    table: Vec<u32>,
6}
7
8impl PrimeTable {
9    pub fn new(max_n: u32) -> Self {
10        let mut table = vec![1; (max_n as usize).div_ceil(2)];
11        table[0] = 0;
12        for i in (3..).step_by(2) {
13            let i2 = i * i;
14            if i2 > max_n {
15                break;
16            }
17            if table[i as usize >> 1] == 1 {
18                for j in (i2..=max_n).step_by(i as usize * 2) {
19                    if table[j as usize >> 1] == 1 {
20                        table[j as usize >> 1] = i;
21                    }
22                }
23            }
24        }
25        PrimeTable { table }
26    }
27    pub fn is_prime(&self, n: u32) -> bool {
28        n == 2 || n % 2 == 1 && self.table[n as usize >> 1] == 1
29    }
30    pub fn primes(&self) -> impl Iterator<Item = u32> + '_ {
31        once(2).chain(self.table.iter().enumerate().filter_map(|(i, b)| {
32            if *b == 1 {
33                Some(i as u32 * 2 + 1)
34            } else {
35                None
36            }
37        }))
38    }
39    pub fn trial_division<F>(&self, mut n: u32, mut f: F)
40    where
41        F: FnMut(u32, u32),
42    {
43        let k = n.trailing_zeros();
44        if k > 0 {
45            f(2, k);
46        }
47        n >>= k;
48        while self.table[n as usize >> 1] > 1 {
49            let p = self.table[n as usize >> 1];
50            let mut cnt = 1;
51            n /= p;
52            while self.table[n as usize >> 1] == p {
53                n /= p;
54                cnt += 1;
55            }
56            if n == p {
57                cnt += 1;
58                n /= p;
59            }
60            f(p, cnt);
61        }
62        if n > 1 {
63            f(n, 1);
64        }
65    }
66    pub fn prime_factors(&self, n: u32) -> Vec<(u32, u32)> {
67        let mut factors = vec![];
68        self.trial_division(n, |p, c| factors.push((p, c)));
69        factors
70    }
71    pub fn count_divisors(&self, n: u32) -> u32 {
72        let mut divisor_cnt = 1;
73        self.trial_division(n, |_, cnt| divisor_cnt *= cnt + 1);
74        divisor_cnt
75    }
76    pub fn divisors(&self, n: u32) -> Vec<u32> {
77        let mut d = vec![1u32];
78        self.trial_division(n, |p, c| {
79            let k = d.len();
80            let mut acc = 1;
81            for _ in 0..c {
82                acc *= p;
83                for i in 0..k {
84                    d.push(d[i] * acc);
85                }
86            }
87        });
88        d.sort_unstable();
89        d
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::tools::Xorshift;
97
98    pub fn divisors(n: u32) -> Vec<u32> {
99        let mut res = vec![];
100        for i in 1..(n as f32).sqrt() as u32 + 1 {
101            if n.is_multiple_of(i) {
102                res.push(i);
103                if i * i != n {
104                    res.push(n / i);
105                }
106            }
107        }
108        res.sort_unstable();
109        res
110    }
111
112    #[test]
113    fn test_prime_table() {
114        for n in (1..=30).chain([1000, 10_000]) {
115            let primes = PrimeTable::new(n);
116            for x in 1..=n {
117                let is_prime = x >= 2 && (2..).take_while(|&d| d * d <= x).all(|d| x % d != 0);
118                assert_eq!(primes.is_prime(x), is_prime);
119                let factors = primes.prime_factors(x);
120                assert_eq!(factors.iter().map(|&(p, e)| p.pow(e)).product::<u32>(), x);
121                assert!(factors.windows(2).all(|w| w[0].0 < w[1].0));
122                for &(p, e) in &factors {
123                    assert!(p >= 2 && e > 0);
124                    assert!((2..).take_while(|&d| d * d <= p).all(|d| p % d != 0));
125                }
126                assert_eq!(
127                    primes.count_divisors(x),
128                    (1..=x).filter(|d| x % d == 0).count() as u32
129                );
130            }
131        }
132    }
133
134    #[test]
135    fn test_divisors() {
136        let mut rng = Xorshift::default();
137        let pt = PrimeTable::new(200001);
138        for n in (1..1000).chain(rng.random_iter(1..=200000).take(100)) {
139            assert_eq!(pt.divisors(n), divisors(n));
140        }
141    }
142}
competitive/math/prime_table.rs

competitive/math/
prime_table.rs
