sfcpl/prime/
sieve.rs

1use cargo_snippet::snippet;
2
3/// 初期化の際にだけエラストテネスの篩を使って素数のリストを生成
4///
5/// `O(n)`
6#[snippet("sieve")]
7pub struct Sieve {
8    size: usize,
9    /// `(0..=n)`までの範囲で
10    /// `is_prime[i]` => `i`は素数
11    is_prime: Vec<bool>,
12    /// `(0..=n)`の全ての素数
13    primes: Vec<usize>,
14}
15
16#[snippet("sieve")]
17impl Sieve {
18    /// 初期化、サイズが重要
19    pub fn new(n: usize) -> Self {
20        let mut spf = vec![None; n + 1];
21        let mut is_prime = vec![true; n + 1];
22        let mut primes = Vec::new();
23
24        is_prime[0] = false;
25        is_prime[1] = false;
26
27        for i in 2..n + 1 {
28            if is_prime[i] {
29                primes.push(i);
30                spf[i] = Some(i);
31            }
32
33            for prime in &primes {
34                if i * prime >= n + 1 || prime > &spf[i].unwrap() {
35                    break;
36                }
37
38                is_prime[i * prime] = false;
39                spf[i * prime] = Some(*prime);
40            }
41        }
42
43        Self {
44            size: n,
45            is_prime,
46            primes,
47        }
48    }
49
50    /// 自分自身の素数リストの有効な範囲
51    pub fn size(&self) -> usize {
52        self.size
53    }
54
55    /// `n`以下の全ての素数のリストを作る
56    pub fn primes(&self, n: usize) -> Vec<usize> {
57        assert!(self.size() >= n);
58        self.primes
59            .iter()
60            .take_while(|x| **x <= n)
61            .cloned()
62            .collect()
63    }
64
65    /// `n`が素数かどうか
66    ///
67    /// # Panic
68    /// `n > self.size`でindexing panicする
69    pub fn is_prime(&self, n: usize) -> bool {
70        self.is_prime[n]
71    }
72}
73
74/// Seiveテーブルを用いた素因数分解
75///
76/// `n`を素因数分解するためには最小で、
77/// `√n`までのサイズの素数テーブルが必要
78///
79/// # Panic
80/// `sieve`のサイズが`√n`未満で不十分な場合にpanicします
81#[snippet("sieve")]
82pub fn factorizations_with_sieve(sieve: &Sieve, mut n: usize) -> Vec<(usize, usize)> {
83    // Panic
84    assert!(sieve.size.pow(2) >= n);
85
86    let mut res = Vec::new();
87    let ps = sieve.primes((n as f64).sqrt().ceil() as usize);
88
89    for p in ps {
90        let mut c = 0usize;
91        while n % p == 0 {
92            n /= p;
93            c += 1;
94        }
95        if c != 0 {
96            res.push((p, c));
97        }
98    }
99
100    // `n`が素数だった
101    if n > 1 {
102        res.push((n, 1));
103    }
104
105    res
106}
107
108#[test]
109fn sieve_test() {
110    let sieve = Sieve::new(10);
111    assert_eq!(vec![2, 3, 5, 7], sieve.primes(10));
112    assert_eq!(sieve.is_prime[2], true);
113    assert_eq!(sieve.is_prime[9], false);
114}
115
116#[test]
117fn fact_test() {
118    let sieve = Sieve::new(100000);
119    let mut f = factorizations_with_sieve(&sieve, 12);
120    f.sort();
121    assert_eq!(f, vec![(2, 2), (3, 1)]);
122    let mut f = factorizations_with_sieve(&sieve, 107);
123    f.sort();
124    assert_eq!(f, vec![(107, 1)]);
125    let mut f = factorizations_with_sieve(&sieve, 1000000007);
126    f.sort();
127    assert_eq!(f, vec![(1000000007, 1)]);
128}
sfcpl/prime/sieve.rs

sfcpl/prime/
sieve.rs
