linear_sieve/
lib.rs

1pub struct LinearSieve {
2    lpf: Vec<usize>,
3    lpf_e: Vec<(usize, u32)>,
4    pr: Vec<usize>,
5}
6
7impl LinearSieve {
8    pub fn new(n: usize) -> Self {
9        let mut lpf = vec![1; n + 1];
10        let mut pr = vec![];
11        for i in 2..=n {
12            if lpf[i] == 1 {
13                lpf[i] = i;
14                pr.push(i);
15            }
16            let lpf_i = lpf[i];
17            for &j in pr.iter().take_while(|&&j| j <= lpf_i.min(n / i)) {
18                lpf[i * j] = j;
19            }
20        }
21        let mut lpf_e = vec![(1, 0); n + 1];
22        for i in 2..=n {
23            let p = lpf[i];
24            let j = i / p;
25            lpf_e[i] = if lpf[j] == p {
26                (lpf_e[j].0 * p, lpf_e[j].1 + 1)
27            } else {
28                (lpf[i], 1)
29            };
30        }
31        Self { lpf, lpf_e, pr }
32    }
33
34    pub fn is_prime(&self, i: usize) -> bool { i >= 2 && self.lpf[i] == i }
35    pub fn lpf(&self, i: usize) -> Option<usize> {
36        (i >= 2).then(|| self.lpf[i])
37    }
38
39    pub fn factors_dup(&self, i: usize) -> impl Iterator<Item = usize> + '_ {
40        std::iter::successors(Some(i), move |&i| Some(i / self.lpf[i]))
41            .take_while(|&i| i > 1)
42            .map(move |i| self.lpf[i])
43    }
44
45    pub fn factors(&self, i: usize) -> impl Iterator<Item = (usize, u32)> + '_ {
46        std::iter::successors(Some(i), move |&i| Some(i / self.lpf_e[i].0))
47            .take_while(|&i| i > 1)
48            .map(move |i| (self.lpf[i], self.lpf_e[i].1))
49    }
50
51    pub fn euler_phi(&self, i: usize) -> usize {
52        std::iter::successors(Some(i), move |&i| Some(i / self.lpf_e[i].0))
53            .take_while(|&i| i > 1)
54            .map(|i| self.lpf_e[i].0 / self.lpf[i] * (self.lpf[i] - 1))
55            .product()
56    }
57
58    pub fn euler_phi_star(&self, i: usize) -> usize {
59        match i {
60            0..=2 => i / 2,
61            _ => 1 + self.euler_phi_star(self.euler_phi(i)),
62        }
63    }
64
65    pub fn divisors(
66        &self,
67        i: usize,
68    ) -> impl Iterator<Item = usize> + DoubleEndedIterator {
69        let mut res = vec![1];
70        for (p, e) in self.factors(i) {
71            let mut tmp = vec![];
72            let mut pp = 1;
73            for _ in 1..=e {
74                pp *= p;
75                tmp.extend(res.iter().map(|&x| x * pp));
76            }
77            res.extend(tmp);
78        }
79        res.sort_unstable();
80        res.into_iter()
81    }
82
83    pub fn divisors_count(&self, i: usize) -> usize {
84        self.factors(i).map(|(_, e)| e as usize + 1).product()
85    }
86
87    pub fn divisors_sum(&self, i: usize) -> usize {
88        std::iter::successors(Some(i), move |&i| Some(i / self.lpf_e[i].0))
89            .take_while(|&i| i > 1)
90            .map(|i| (self.lpf_e[i].0 * self.lpf[i] - 1) / (self.lpf[i] - 1))
91            .product()
92    }
93
94    pub fn primes(
95        &self,
96    ) -> impl Iterator<Item = usize> + DoubleEndedIterator + '_ {
97        self.pr.iter().copied()
98    }
99
100    pub fn dp<T>(
101        &self,
102        zero: T,
103        one: T,
104        eq: impl Fn(&T, usize) -> T,
105        gt: impl Fn(&T, usize) -> T,
106    ) -> Vec<T> {
107        let n = self.lpf.len() - 1;
108
109        let mut res = vec![zero, one];
110        if n <= 1 {
111            res.truncate(n + 1);
112            return res;
113        }
114
115        res.reserve(n + 1);
116        for i in 2..=n {
117            let lpf = self.lpf[i];
118            let j = i / lpf;
119            let prev = &res[j];
120            let cur =
121                if lpf == self.lpf[j] { eq(prev, lpf) } else { gt(prev, lpf) };
122            res.push(cur);
123        }
124        res
125    }
126}
127
128#[test]
129fn sanity_check() {
130    let ls = LinearSieve::new(60);
131
132    let primes =
133        vec![2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59];
134
135    assert!(ls.primes().eq(primes.iter().copied()));
136    assert!((1..=60).all(|i| primes.contains(&i) == ls.is_prime(i)));
137
138    assert_eq!(ls.lpf(1), None);
139    assert_eq!(ls.lpf(3), Some(3));
140    assert_eq!(ls.lpf(24), Some(2));
141
142    assert!(ls.factors_dup(1).eq(None));
143    assert!(ls.factors_dup(60).eq([2, 2, 3, 5]));
144
145    assert!(ls.factors(1).eq(None));
146    assert!(ls.factors(60).eq([(2, 2), (3, 1), (5, 1)]));
147
148    assert_eq!(ls.euler_phi(1), 1);
149    assert_eq!(ls.euler_phi(35), 24);
150    assert_eq!(ls.euler_phi(60), 16);
151
152    assert_eq!(ls.euler_phi_star(1), 0);
153    assert_eq!(ls.euler_phi_star(35), 5);
154    assert_eq!(ls.euler_phi_star(60), 5);
155
156    assert!(ls.divisors(1).eq([1]));
157    assert!(ls.divisors(60).eq([1, 2, 3, 4, 5, 6, 10, 12, 15, 20, 30, 60]));
158
159    assert_eq!(ls.divisors_count(1), 1);
160    assert_eq!(ls.divisors_count(60), 12);
161
162    assert_eq!(ls.divisors_sum(1), 1);
163    assert_eq!(ls.divisors_sum(60), 168);
164}
165
166#[test]
167fn dp() {
168    let ls = LinearSieve::new(10);
169
170    // Moebius mu
171    let mu = ls.dp(0, 1, |_, _| 0, |&x, _| -x);
172    assert_eq!(mu, [0, 1, -1, -1, 0, -1, 1, -1, 0, 0, 1]);
173
174    // Euler phi
175    let phi = ls.dp(0, 1, |&x, p| x * p, |&x, p| x * (p - 1));
176    assert_eq!(phi, [0, 1, 1, 2, 2, 4, 2, 6, 4, 6, 4]);
177
178    // # of distinct prime factors
179    let omega = ls.dp(0, 0, |&x, _| x, |&x, _| x + 1);
180    assert_eq!(omega, [0, 0, 1, 1, 1, 1, 2, 1, 1, 1, 2]);
181
182    // # of prime factors
183    let cap_omega = ls.dp(0, 0, |&x, _| x + 1, |&x, _| x + 1);
184    assert_eq!(cap_omega, [0, 0, 1, 1, 2, 1, 2, 1, 3, 2, 2]);
185
186    // sum of divisors
187    let eq = |&(prod, sum, pow): &_, p| (prod, sum + pow * p, pow * p);
188    let gt = |&(prod, sum, _): &_, p| (prod * sum, 1 + p, p);
189    let sigma: [_; 11] =
190        ls.dp((0, 0, 0), (1, 1, 1), eq, gt).try_into().unwrap();
191    let sigma = sigma.map(|(prod, sum, _)| prod * sum);
192    assert_eq!(sigma, [0, 1, 3, 4, 7, 6, 12, 8, 15, 13, 18]);
193}
linear_sieve/lib.rs

linear_sieve/
lib.rs
