atcoder/lib/math/
sieve.rs

1//! 素数と素因数分解に関するアルゴリズムです。
2
3/// エラトステネスの篩で素数と最小素因数を前計算します。
4///
5/// `0..=limit` の範囲で素数判定と素因数分解を行うため、各整数の最小素因数を保持します。
6/// `is_prime` は `O(1)`、`prime_factorize` は素因数の重複度（指数の合計）に比例する
7/// `O(log n)` 時間で計算できます。素数列挙は `primes()` から行えます。
8///
9/// 構築の計算量は `O(N log log N)`、使用メモリは `O(N)` です。
10///
11/// # Examples
12///
13/// ```
14/// use atcoder::math::sieve::Sieve;
15///
16/// let sieve = Sieve::new(100);
17///
18/// assert!(sieve.is_prime(97));
19/// assert!(!sieve.is_prime(1));
20///
21/// assert_eq!(
22///     sieve.prime_factorize(84),
23///     vec![(2, 2), (3, 1), (7, 1)],
24/// );
25/// ```
26#[derive(Debug, Clone)]
27pub struct Sieve {
28    smallest_prime_factor: Vec<usize>,
29}
30
31impl Sieve {
32    /// `0..=limit` の各整数について最小素因数を前計算します。判定・分解可能な値はこの範囲内です。
33    ///
34    /// 計算量は `O(limit log log limit)`、使用メモリは `O(limit)` です。
35    pub fn new(limit: usize) -> Self {
36        let mut smallest_prime_factor: Vec<usize> = (0..=limit).collect();
37
38        for prime in 2..=limit {
39            // 最小素因数が自分自身なら prime は素数である。
40            if smallest_prime_factor[prime] != prime {
41                continue;
42            }
43
44            // prime * prime のオーバーフローを避ける。
45            if prime > limit / prime {
46                continue;
47            }
48
49            // prime 未満の倍数は、より小さい素因数によって
50            // すでに処理されているため prime^2 から始める。
51            for multiple in (prime * prime..=limit).step_by(prime) {
52                if smallest_prime_factor[multiple] == multiple {
53                    smallest_prime_factor[multiple] = prime;
54                }
55            }
56        }
57
58        Self {
59            smallest_prime_factor,
60        }
61    }
62
63    /// 前計算した範囲の上端 `limit` を返します（有効範囲は `0..=limit()`）。
64    pub fn limit(&self) -> usize {
65        self.smallest_prime_factor.len() - 1
66    }
67
68    /// 前計算範囲内の `n` が素数かどうかを返します。`0` と `1` は `false` です。
69    ///
70    /// 計算量は `O(1)` です。
71    ///
72    /// # Panics
73    ///
74    /// `n > limit` の場合は panic します。
75    pub fn is_prime(&self, n: usize) -> bool {
76        self.assert_in_range(n);
77
78        n >= 2 && self.smallest_prime_factor[n] == n
79    }
80
81    /// `0..=limit()` に含まれる素数を昇順に列挙します。
82    ///
83    /// 全要素を列挙する計算量は `O(limit())` です。イテレータは `self` を借用するため、
84    /// `Sieve` より長く保持することはできません。
85    pub fn primes(&self) -> impl Iterator<Item = usize> + '_ {
86        (2..=self.limit()).filter(|&n| self.smallest_prime_factor[n] == n)
87    }
88
89    /// 前計算範囲内の `n` を素因数分解し、`(素因数, 指数)` の組を素因数の昇順で返します。
90    ///
91    /// `n = 1` の場合は空の `Vec` を返します。`n = 0` は素因数分解できません。
92    /// 計算量は素因数の重複度を含めて `O(log n)` です。
93    ///
94    /// # Panics
95    ///
96    /// - `n = 0` の場合
97    /// - `n > limit` の場合
98    pub fn prime_factorize(&self, mut n: usize) -> Vec<(usize, usize)> {
99        assert!(n > 0, "zero does not have a prime factorization");
100
101        self.assert_in_range(n);
102
103        let mut factors = Vec::new();
104
105        while n > 1 {
106            let prime = self.smallest_prime_factor[n];
107            let mut exponent = 0;
108
109            while n % prime == 0 {
110                n /= prime;
111                exponent += 1;
112            }
113
114            factors.push((prime, exponent));
115        }
116
117        factors
118    }
119
120    fn assert_in_range(&self, n: usize) {
121        assert!(
122            n <= self.limit(),
123            "n exceeds the precomputed range: n = {n}, limit = {}",
124            self.limit(),
125        );
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::Sieve;
132
133    #[test]
134    fn identifies_primes() {
135        let sieve = Sieve::new(20);
136
137        assert!(!sieve.is_prime(0));
138        assert!(!sieve.is_prime(1));
139        assert!(sieve.is_prime(2));
140        assert!(sieve.is_prime(19));
141        assert!(!sieve.is_prime(20));
142
143        assert_eq!(
144            sieve.primes().collect::<Vec<_>>(),
145            vec![2, 3, 5, 7, 11, 13, 17, 19],
146        );
147    }
148
149    #[test]
150    fn factorizes_numbers() {
151        let sieve = Sieve::new(100);
152
153        assert_eq!(sieve.prime_factorize(1), vec![]);
154        assert_eq!(sieve.prime_factorize(2), vec![(2, 1)]);
155
156        assert_eq!(sieve.prime_factorize(84), vec![(2, 2), (3, 1), (7, 1)],);
157
158        assert_eq!(sieve.prime_factorize(97), vec![(97, 1)],);
159    }
160}
atcoder/lib/math/sieve.rs

atcoder/lib/math/
sieve.rs
