competitive_hpp/prime/
eratosthenes.rs

1use num::traits::PrimInt;
2use std::collections::HashMap;
3#[derive(Eq, PartialEq, Clone, Debug)]
4pub struct Eratosthenes {
5    primes: Vec<usize>,
6    flags: Vec<usize>,
7}
8
9/// # Eratosthenes
10///
11/// Example:
12/// ```
13/// use competitive_hpp::prelude::*;
14/// let eratosthenes = Eratosthenes::new(100);
15///
16/// assert!(!eratosthenes.is_prime(-1));
17/// assert!(!eratosthenes.is_prime(0));
18/// assert!(eratosthenes.is_prime(2));
19/// assert!(!eratosthenes.is_prime(10));
20/// assert!(eratosthenes.is_prime(11));
21///
22/// let mut map = HashMap::new();
23/// map.insert(2, 1);
24/// map.insert(5, 2);
25/// assert_eq!(map, eratosthenes.factorization(50));
26/// ```
27#[allow(clippy::many_single_char_names)]
28impl Eratosthenes {
29    pub fn new<T>(n: T) -> Self
30    where
31        T: PrimInt,
32    {
33        let n = n.to_usize().unwrap();
34        let bits = 32;
35        let flags_num = n / bits + 1;
36
37        let defaults: Vec<usize> = vec![0x5D75D75D, 0xD75D75D7, 0x75D75D75];
38
39        let (mut i, mut f, mut j);
40        let max = ((n as f64).sqrt() as usize) + 1;
41
42        let mut flags: Vec<usize> = (0..flags_num).map(|i| defaults[i % 3]).collect();
43        flags[flags_num - 1] &= (1 << (n % bits + 1)) - 1;
44
45        i = 5;
46        f = 4;
47        while i <= max {
48            let t = (flags[i / bits] >> (i % bits)) & 1 == 1;
49            if !t {
50                j = i * (i | 1);
51                while j <= n {
52                    flags[j / bits] |= 1 << (j % bits);
53                    j += i * 2;
54                }
55            }
56            f = 6 - f;
57            i += f;
58        }
59
60        flags[0] &= !0b1100;
61        flags[0] |= 0b11;
62
63        let mut primes = vec![];
64        for p in 2..n + 1 {
65            if (flags[p / bits] >> (p % bits)) & 1 == 0 {
66                primes.push(p);
67            }
68        }
69
70        Eratosthenes { primes, flags }
71    }
72
73    pub fn is_prime<T>(&self, m: T) -> bool
74    where
75        T: PrimInt,
76    {
77        match m.to_usize() {
78            Some(um) => self.flags[um / 32] >> (um % 32) & 1 == 0,
79            None => false,
80        }
81    }
82
83    pub fn factorization<T>(&self, n: T) -> HashMap<usize, usize>
84    where
85        T: PrimInt,
86    {
87        let mut n = n.to_usize().unwrap();
88        let mut factors: HashMap<usize, usize> = HashMap::new();
89        for &p in &self.primes {
90            while n % p == 0 {
91                n /= p;
92                factors
93                    .entry(p)
94                    .and_modify(|entry| *entry += 1)
95                    .or_insert(1);
96            }
97            if p > n {
98                break;
99            }
100        }
101        factors
102    }
103}
104
105#[cfg(test)]
106mod test {
107    use super::*;
108
109    #[test]
110    fn test_eratosthenes() {
111        let eratosthenes = Eratosthenes::new(100);
112
113        assert!(!eratosthenes.is_prime(-1));
114        assert!(!eratosthenes.is_prime(0));
115        assert!(!eratosthenes.is_prime(1));
116        assert!(eratosthenes.is_prime(2));
117        assert!(eratosthenes.is_prime(3));
118        assert!(!eratosthenes.is_prime(4));
119        assert!(eratosthenes.is_prime(5));
120        assert!(!eratosthenes.is_prime(6));
121        assert!(eratosthenes.is_prime(7));
122        assert!(!eratosthenes.is_prime(8));
123        assert!(!eratosthenes.is_prime(9));
124        assert!(!eratosthenes.is_prime(10));
125        assert!(eratosthenes.is_prime(11));
126
127        let mut map = HashMap::new();
128        map.insert(2, 1);
129        map.insert(5, 2);
130
131        assert_eq!(map, eratosthenes.factorization(50));
132        assert_ne!(map, eratosthenes.factorization(100));
133    }
134}
competitive_hpp/prime/eratosthenes.rs

competitive_hpp/prime/
eratosthenes.rs
