competitive_library/algorithm/
fast_eratosthenes.rs

1//! エラトステネス
2
3///エラトステネスの篩
4pub struct Eratosthenes {
5    flags_: Vec<u8>,
6    n: usize,
7}
8impl Eratosthenes {
9    const K_MASK: [[u8; 8]; 8] = [
10        [0xfe, 0xfd, 0xfb, 0xf7, 0xef, 0xdf, 0xbf, 0x7f],
11        [0xfd, 0xdf, 0xef, 0xfe, 0x7f, 0xf7, 0xfb, 0xbf],
12        [0xfb, 0xef, 0xfe, 0xbf, 0xfd, 0x7f, 0xf7, 0xdf],
13        [0xf7, 0xfe, 0xbf, 0xdf, 0xfb, 0xfd, 0x7f, 0xef],
14        [0xef, 0x7f, 0xfd, 0xfb, 0xdf, 0xbf, 0xfe, 0xf7],
15        [0xdf, 0xf7, 0x7f, 0xfd, 0xbf, 0xfe, 0xef, 0xfb],
16        [0xbf, 0xfb, 0xf7, 0x7f, 0xfe, 0xef, 0xdf, 0xfd],
17        [0x7f, 0xbf, 0xdf, 0xef, 0xf7, 0xfb, 0xfd, 0xfe],
18    ];
19
20    const C0: [[usize; 8]; 8] = [
21        [0, 0, 0, 0, 0, 0, 0, 1],
22        [1, 1, 1, 0, 1, 1, 1, 1],
23        [2, 2, 0, 2, 0, 2, 2, 1],
24        [3, 1, 1, 2, 1, 1, 3, 1],
25        [3, 3, 1, 2, 1, 3, 3, 1],
26        [4, 2, 2, 2, 2, 2, 4, 1],
27        [5, 3, 1, 4, 1, 3, 5, 1],
28        [6, 4, 2, 4, 2, 4, 6, 1],
29    ];
30    const MOD_30: [usize; 8] = [1, 7, 11, 13, 17, 19, 23, 29];
31    const C1: [usize; 8] = [6, 4, 2, 4, 2, 4, 6, 2];
32
33    ///初期化
34    ///素数フラグを処理
35    ///- param n:usize 探索上限
36    pub fn new(n: usize) -> Self {
37        if n > 10_000_000_000 {
38            panic!();
39        }
40
41        let size = n / 30 + usize::from(n % 30 != 0);
42        let mut flags_ = vec![0xff_u8; size];
43        flags_[0] = 0xfe;
44
45        let remainder = n % 30;
46        flags_[size - 1] = match remainder {
47            1 => 0x0,
48            2..=7 => 0x1,
49            8..=11 => 0x3,
50            12..=13 => 0x7,
51            14..=17 => 0xf,
52            18..=19 => 0x1f,
53            20..=23 => 0x3f,
54            // 24..=29
55            _ => 0x7f,
56        };
57
58        let quart_x = ((n as f64).sqrt() + 1.0) as usize / 30 + 1;
59
60        for i in 0..quart_x {
61            let mut flags: u8 = flags_[i];
62
63            while flags != 0 {
64                let lsb = flags & flags.wrapping_neg();
65                let i_bit = lsb.trailing_zeros() as usize;
66
67                let m = Self::MOD_30[i_bit];
68
69                let mut k = i_bit;
70                let mut j = i * (30 * i + 2 * m) + (m * m) / 30;
71
72                while j < flags_.len() {
73                    flags_[j] &= Self::K_MASK[i_bit][k];
74
75                    j += i * Self::C1[k] + Self::C0[i_bit][k];
76                    k = (k + 1) & 7;
77                }
78                flags &= flags - 1;
79            }
80        }
81
82        Self { flags_, n }
83    }
84
85    ///素数の個数をカウント
86    pub fn count(&mut self) -> usize {
87        let mut ret = [2usize, 3, 5].iter().take_while(|x| self.n >= **x).count(); // count 2, 3, 5
88        for f in &self.flags_ {
89            ret += f.count_ones() as usize;
90        }
91        ret
92    }
93
94    ///フラグから素数配列を生成
95    pub fn primes(&self) -> Vec<usize> {
96        let mut ret = vec![];
97
98        [2usize, 3, 5]
99            .iter()
100            .take_while(|x| self.n >= **x)
101            .for_each(|x| ret.push(*x));
102
103        for (i, f) in self.flags_.iter().enumerate() {
104            for (ii, m) in Self::MOD_30.iter().enumerate() {
105                if (*f & (1 << ii)) != 0 {
106                    ret.push(30 * i + *m);
107                }
108            }
109        }
110        ret
111    }
112}
113
114#[cfg(test)]
115mod tests {
116
117    use super::*;
118
119    #[test]
120    fn test_era_1e9() {
121        let mut e = Eratosthenes::new(100_000_000);
122
123        assert_eq!(e.count(), 5_761_455);
124    }
125    #[test]
126    fn test_era_zero() {
127        let mut e = Eratosthenes::new(1);
128
129        assert_eq!(e.count(), 0);
130    }
131}