competitive_library/algorithm/
eratosthenes.rs

1//! エラトステネス
2
3use std::collections::BTreeMap;
4
5pub struct Sieve {
6    n: usize,
7    pub table: Vec<i64>,
8    pub primes: Vec<usize>,
9}
10impl Sieve {
11    pub fn new(n: usize) -> Self {
12        let mut primes = vec![];
13        let mut table = vec![0_i64; n + 1];
14        for i in 2..n + 1 {
15            if table[i] == 0 {
16                primes.push(i);
17                table[i] = i as i64;
18            }
19            for &p in &primes {
20                if p * i > n {
21                    break;
22                }
23                table[p * i] = p as i64;
24            }
25        }
26        Sieve { n, table, primes }
27    }
28
29    pub fn factorize(&self, n: i64) -> BTreeMap<i64, i64> {
30        assert!(self.n as i64 >= n);
31
32        let mut map = BTreeMap::new();
33        let mut target = n;
34
35        while target > 1 {
36            let p = self.table[target as usize];
37            let mut count = 0;
38
39            while self.table[target as usize] == p {
40                target /= p;
41                count += 1;
42            }
43            map.insert(p, count);
44        }
45        map
46    }
47
48    // 約数列挙
49    pub fn divisors(&self, n: i64) -> Vec<i64> {
50        assert_ne!(n, 0);
51
52        let mut ret = vec![1];
53        let factor = self.factorize(n);
54
55        for (k, exp) in factor {
56            for i in 0..ret.len() {
57                let mut v = 1;
58                for _ in 0..exp {
59                    v *= k;
60                    ret.push(ret[i] * v);
61                }
62            }
63        }
64        ret
65    }
66
67    pub fn is_prime(&self, n: i64) -> bool {
68        assert!(self.n as i64 >= n);
69        self.table[n as usize] == n
70    }
71}
72
73#[cfg(test)]
74mod tests {
75
76    use super::*;
77
78    #[test]
79    fn test_prime() {
80        assert_eq!(Sieve::new(1_000_000).primes.len(), 78_498);
81    }
82
83    #[test]
84    fn test_factorize() {
85        let ls = Sieve::new(1_000_000);
86        let case = vec![
87            31, 4657, 3, 65732, 7836, 1278, 8615, 798_179, 425, 38715, 3272, 149, 7, 48, 97, 5823,
88            517, 8231, 6986, 751, 8651, 671, 7, 23,
89        ];
90
91        for v in case {
92            let f = ls.factorize(v);
93
94            let mut n = 1;
95            for (k, value) in f {
96                n *= k.pow(value as u32);
97            }
98            assert_eq!(v, n);
99        }
100    }
101
102    use std::collections::HashMap;
103    #[test]
104    fn test_divide() {
105        let ls = Sieve::new(1_000_000);
106        let map = {
107            let mut ret = HashMap::new();
108            ret.insert(1, vec![1]);
109            ret.insert(2, vec![1, 2]);
110            ret.insert(3, vec![1, 3]);
111            ret.insert(4, vec![1, 2, 4]);
112            ret.insert(6, vec![1, 2, 3, 6]);
113            ret.insert(20, vec![1, 2, 4, 5, 10, 20]);
114            ret.insert(25, vec![1, 5, 25]);
115            ret.insert(30, vec![1, 2, 3, 5, 6, 10, 15, 30]);
116            ret.insert(
117                2520,
118                vec![
119                    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 15, 18, 20, 21, 24, 28, 30, 35, 36, 40,
120                    42, 45, 56, 60, 63, 70, 72, 84, 90, 105, 120, 126, 140, 168, 180, 210, 252,
121                    280, 315, 360, 420, 504, 630, 840, 1260, 2520,
122                ],
123            );
124            ret
125        };
126
127        for (k, v) in map {
128            let mut a = ls.divisors(k);
129            a.sort_unstable();
130            assert_eq!(a, v);
131        }
132    }
133}
competitive_library/algorithm/eratosthenes.rs

competitive_library/algorithm/
eratosthenes.rs
