competitive_library/algorithm/
atkin.rs

1//! アトキンの篩
2pub fn generate_primes(n: u64) -> Vec<bool> {
3    let mut is_prime = vec![false; n as usize + 1];
4    let sqrt_n = (n as f64).sqrt() as u64 + 1;
5
6    for i in 1..sqrt_n {
7        for j in 1..sqrt_n {
8            let ii = i.pow(2);
9            let jj = j.pow(2);
10
11            let mut buff = 4 * ii + jj;
12            let buff_mod12 = buff % 12;
13            if buff <= n && (buff_mod12 == 1 || buff_mod12 == 5) {
14                is_prime[buff as usize] ^= true;
15            }
16
17            buff = 3 * ii + jj;
18            if buff <= n && buff % 12 == 7 {
19                is_prime[buff as usize] ^= true;
20            }
21
22            if i <= j {
23                continue;
24            }
25
26            buff = 3 * ii - jj;
27            if i > j && buff <= n && buff % 12 == 11 {
28                is_prime[buff as usize] ^= true;
29            }
30        }
31    }
32    for i in 5..sqrt_n {
33        if !is_prime[i as usize] {
34            continue;
35        }
36        let k = i * i;
37        for j in (k..n).step_by(k as usize) {
38            is_prime[j as usize] = false;
39        }
40    }
41    is_prime[2] = true;
42    is_prime[3] = true;
43    is_prime
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49    #[test]
50    fn test_atkin() {
51        let prime = generate_primes(1_000_000);
52
53        let count: Vec<_> = prime
54            .iter()
55            .enumerate()
56            .filter(|x| *x.1)
57            .map(|x| x.0)
58            .collect();
59        assert_eq!(count.len(), 78498);
60        assert_eq!(count[0], 2);
61    }
62}