nekolib/math/
sieve_n2_plus_1.rs

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

nekolib/math/
sieve_n2_plus_1.rs
