haar_lib/math/prime_test/
segmented.rs

1//! 区間篩
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc412/tasks/abc412_e>
5
6pub use crate::math::prime_test::CheckPrime;
7
8/// 区間篩
9pub struct SegmentedSieve {
10    l: usize,
11    r: usize,
12    data: Vec<bool>,
13}
14
15impl SegmentedSieve {
16    /// `[l, r]`の区間篩を作る。
17    ///
18    /// `check`は$\sqrt{r}$の以下の素数判定が可能であること。
19    pub fn new(l: usize, r: usize, check: &impl CheckPrime<usize>) -> Self {
20        assert!(l <= r);
21
22        let d = r - l + 1;
23        let primes = (2..)
24            .take_while(|i| i * i <= r)
25            .filter(|&i| check.is_prime(i))
26            .collect::<Vec<_>>();
27        let mut data = vec![true; d];
28
29        for p in primes {
30            let mut from = (l + p - 1) / p * p;
31            if p == from {
32                from = p * 2;
33            }
34
35            for i in (from..=r).step_by(p) {
36                data[i - l] = false;
37            }
38        }
39
40        Self { l, r, data }
41    }
42}
43
44impl CheckPrime<usize> for SegmentedSieve {
45    fn is_prime(&self, value: usize) -> bool {
46        assert!(self.l <= value && value <= self.r);
47        self.data[value - self.l]
48    }
49}