1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
//! $n^2+1$ 型素数の篩。

use std::fmt::Debug;

/// $n^2+1$ 型素数の篩。
///
/// # Idea
/// $n^2+1$ が素因数 $p < n$ を持つとき、以下が成り立つ。
/// $$ n^2+1 \\equiv 0 \\pmod{p}. $$
/// このとき、各 $k$ に対して以下が成り立つ。
///
/// - $(n+kp)^2+1 \\equiv 0 \\pmod{p}$,
/// - $(kp-n)^2+1 \\equiv 0 \\pmod{p}$.
///
/// # References
/// - <https://inamori.hateblo.jp/entry/20100510/p1>
///
/// # See also
/// - <https://inamori.hateblo.jp/entry/20110930/p1>
///     - $n^2+n+1$ 型素数についての話。
/// - <https://twitter.com/min_25_/status/1418781997892136960>
/// - <https://twitter.com/min_25_/status/1418794801625853952>
///     - $an^2+bn+c$ 型素数についての話。
#[derive(Clone, Debug)]
pub struct SieveN2Plus1 {
    a: Vec<usize>,
    f: Vec<Vec<(usize, u32)>>,
}

impl SieveN2Plus1 {
    /// 初期化する。
    pub fn new(n: usize) -> Self {
        let mut a: Vec<_> = (0..=n).map(|k| k * k + 1).collect();
        let mut f = vec![vec![]; n + 1];
        f[1] = vec![(2, 1)];
        for j in (3..=n).step_by(2) {
            let (d, e) = Self::div_pow(a[j], 2);
            f[j].push((2, e));
            a[j] = d;
        }
        for i in 2..=n {
            let p = a[i];
            if p == 1 {
                continue;
            }
            if p == i * i + 1 {
                f[i].push((p, 1));
            }
            let init1 = if p == i * i + 1 { p + i } else { i };
            let init2 = (n + p - 1) / p * p;
            for j in (init1..=n).step_by(p).chain((init2..=n).step_by(p)) {
                let (d, e) = Self::div_pow(a[j], p);
                if e > 0 {
                    f[j].push((p, e));
                    a[j] = d;
                }
            }
        }
        Self { a, f }
    }

    fn div_pow(mut a: usize, p: usize) -> (usize, u32) {
        let mut e = 0;
        while a % p == 0 {
            a /= p;
            e += 1;
        }
        (a, e)
    }

    /// $n^2+1$ の形の素数を返す。
    ///
    /// # Examples
    /// ```
    /// use nekolib::math::SieveN2Plus1;
    ///
    /// let ss = SieveN2Plus1::new(10);
    /// let primes: Vec<_> = ss.primes().collect();
    /// assert_eq!(primes, [2, 5, 17, 37, 101]);
    /// ```
    pub fn primes(&self) -> impl Iterator<Item = usize> + '_ {
        (1..self.a.len()).filter(move |&i| self.a[i] != 1).map(|i| i * i + 1)
    }

    /// $n^2+1$ が素数のとき真を返す。
    ///
    /// # Examples
    /// ```
    /// use nekolib::math::SieveN2Plus1;
    ///
    /// let ss = SieveN2Plus1::new(10);
    /// assert!(ss.is_prime(2));  // 5 is prime
    /// assert!(!ss.is_prime(5));  // 26 = 2 * 13
    /// ```
    pub fn is_prime(&self, n: usize) -> bool { n > 0 && self.a[n] != 1 }

    /// $n^2+1$ を素因数分解する。
    ///
    /// 底の昇順とは限らないので注意。最小の反例は `n = 21`。
    ///
    /// # Examples
    /// ```
    /// use nekolib::math::SieveN2Plus1;
    ///
    /// let ss = SieveN2Plus1::new(10);
    /// assert_eq!(ss.factors(0).next(), None);
    /// assert_eq!(ss.factors(4).collect::<Vec<_>>(), [(17, 1)]);
    /// assert_eq!(ss.factors(7).collect::<Vec<_>>(), [(2, 1), (5, 2)]);
    /// ```
    pub fn factors(&self, n: usize) -> impl Iterator<Item = (usize, u32)> + '_ {
        self.f[n].iter().cloned()
    }

    /// $n^2+1$ を素因数を列挙する。重複あり。
    ///
    /// 底の昇順とは限らないので注意。最小の反例は `n = 21`。
    ///
    /// # Examples
    /// ```
    /// use nekolib::math::SieveN2Plus1;
    ///
    /// let ss = SieveN2Plus1::new(10);
    /// assert_eq!(ss.factors_dup(0).next(), None);
    /// assert_eq!(ss.factors_dup(4).collect::<Vec<_>>(), [17]);
    /// assert_eq!(ss.factors_dup(7).collect::<Vec<_>>(), [2, 5, 5]);
    /// ```
    pub fn factors_dup(&self, n: usize) -> impl Iterator<Item = usize> + '_ {
        self.factors(n).flat_map(|(p, e)| std::iter::repeat(p).take(e as usize))
    }
}

#[test]
fn test() {
    use linear_sieve::LinearSieve;

    let n = 3000;
    let ls = LinearSieve::new(n * n + 1);
    let ss = SieveN2Plus1::new(n);
    for i in 0..=n {
        assert_eq!(ls.is_prime(i * i + 1), ss.is_prime(i));

        {
            // factors
            let expected: Vec<_> = ls.factors_dup(i * i + 1).collect();
            let mut actual: Vec<_> = ss.factors_dup(i).collect();
            actual.sort_unstable();
            assert_eq!(actual, expected);
        }
    }
}