nekolib/math/
sieve_n2_plus_n_plus_1.rs

1//! $n^2+n+1$ 型素数の篩。
2
3use std::fmt::Debug;
4
5/// $n^2+n+1$ 型素数の篩。
6///
7/// # Idea
8/// `todo!()`
9///
10/// # References
11/// - <https://inamori.hateblo.jp/entry/20110930/p1>
12#[derive(Clone, Debug)]
13pub struct SieveN2PlusNPlus1 {
14    a: Vec<usize>,
15    f: Vec<Vec<(usize, u32)>>,
16}
17
18impl SieveN2PlusNPlus1 {
19    /// 初期化する。
20    pub fn new(n: usize) -> Self {
21        let mut a: Vec<_> = (0..=n).map(|k| k * (k + 1) + 1).collect();
22        let mut f = vec![vec![]; n + 1];
23        f[1] = vec![(3, 1)];
24        for j in (4..=n).step_by(3) {
25            let (d, e) = Self::div_pow(a[j], 3);
26            f[j].push((3, e));
27            a[j] = d;
28        }
29        for i in 2..=n {
30            let p = a[i];
31            if p == 1 {
32                continue;
33            }
34            if p == i * (i + 1) + 1 {
35                f[i].push((p, 1));
36            }
37            let init1 = if p == i * (i + 1) + 1 { p + i } else { i };
38            let init2 = (n + p - 1) / p * p;
39            for j in (init1..=n).step_by(p).chain((init2..=n).step_by(p)) {
40                let (d, e) = Self::div_pow(a[j], p);
41                if e > 0 {
42                    f[j].push((p, e));
43                    a[j] = d;
44                }
45            }
46        }
47        Self { a, f }
48    }
49
50    fn div_pow(mut a: usize, p: usize) -> (usize, u32) {
51        let mut e = 0;
52        while a % p == 0 {
53            a /= p;
54            e += 1;
55        }
56        (a, e)
57    }
58
59    /// $n^2+n+1$ の形の素数を返す。
60    ///
61    /// # Examples
62    /// ```
63    /// use nekolib::math::SieveN2PlusNPlus1;
64    ///
65    /// let ss = SieveN2PlusNPlus1::new(10);
66    /// let primes: Vec<_> = ss.primes().collect();
67    /// assert_eq!(primes, [3, 7, 13, 31, 43, 73]);
68    /// ```
69    pub fn primes(&self) -> impl Iterator<Item = usize> + '_ {
70        (1..self.a.len())
71            .filter(move |&i| self.a[i] == i * (i + 1) + 1)
72            .map(|i| i * (i + 1) + 1)
73    }
74
75    /// $n^2+n+1$ が素数のとき真を返す。
76    ///
77    /// # Examples
78    /// ```
79    /// use nekolib::math::SieveN2PlusNPlus1;
80    ///
81    /// let ss = SieveN2PlusNPlus1::new(10);
82    /// assert!(ss.is_prime(3));  // 13 is prime
83    /// assert!(!ss.is_prime(4));  // 21 = 3 * 7
84    /// ```
85    pub fn is_prime(&self, n: usize) -> bool {
86        n > 0 && self.a[n] == n * (n + 1) + 1
87    }
88
89    /// $n^2+n+1$ を素因数分解する。
90    ///
91    /// 底の昇順とは限らないので注意。
92    ///
93    /// # Examples
94    /// ```
95    /// use nekolib::math::SieveN2PlusNPlus1;
96    ///
97    /// let ss = SieveN2PlusNPlus1::new(20);
98    /// assert_eq!(ss.factors(0).next(), None);
99    /// assert_eq!(ss.factors(7).collect::<Vec<_>>(), [(3, 1), (19, 1)]);
100    /// assert_eq!(ss.factors(18).collect::<Vec<_>>(), [(7, 3)]);
101    /// ```
102    pub fn factors(&self, n: usize) -> impl Iterator<Item = (usize, u32)> + '_ {
103        self.f[n].iter().cloned()
104    }
105
106    /// $n^2+1$ を素因数を列挙する。重複あり。
107    ///
108    /// 底の昇順とは限らないので注意。
109    ///
110    /// # Examples
111    /// ```
112    /// use nekolib::math::SieveN2PlusNPlus1;
113    ///
114    /// let ss = SieveN2PlusNPlus1::new(20);
115    /// assert_eq!(ss.factors_dup(0).next(), None);
116    /// assert_eq!(ss.factors_dup(6).collect::<Vec<_>>(), [43]);
117    /// assert_eq!(ss.factors_dup(18).collect::<Vec<_>>(), [7, 7, 7]);
118    /// ```
119    pub fn factors_dup(&self, n: usize) -> impl Iterator<Item = usize> + '_ {
120        self.factors(n).flat_map(|(p, e)| std::iter::repeat(p).take(e as usize))
121    }
122}
123
124#[test]
125fn test() {
126    use linear_sieve::LinearSieve;
127
128    let n = 3000;
129    let ls = LinearSieve::new(n * (n + 1) + 1);
130    let ss = SieveN2PlusNPlus1::new(n);
131    for i in 0..=n {
132        assert_eq!(ls.is_prime(i * (i + 1) + 1), ss.is_prime(i));
133
134        {
135            // factors
136            let expected: Vec<_> = ls.factors_dup(i * (i + 1) + 1).collect();
137            let mut actual: Vec<_> = ss.factors_dup(i).collect();
138            actual.sort_unstable();
139            assert_eq!(actual, expected);
140        }
141    }
142}
nekolib/math/sieve_n2_plus_n_plus_1.rs

nekolib/math/
sieve_n2_plus_n_plus_1.rs
