nekolib/math/
segmented_factor_sieve.rs

1use super::sqrt;
2use std::fmt::Debug;
3use std::ops::RangeInclusive;
4
5use sqrt::Sqrt;
6
7#[derive(Debug)]
8pub struct SegmentedFactorSieve {
9    ofs: usize,
10    small_p: Vec<usize>,
11    large_p: Vec<Vec<usize>>,
12    large_prod: Vec<usize>,
13}
14
15impl SegmentedFactorSieve {
16    pub fn new(range: RangeInclusive<usize>) -> Self {
17        let &l = range.start();
18        let r = *range.end();
19        let r0 = r.sqrt();
20        let mut large = vec![true; r - l + 1];
21        let mut small_p: Vec<_> = (0..=r0).collect();
22        let mut large_p = vec![vec![]; r - l + 1];
23        let mut large_prod: Vec<_> = (l..=r).collect();
24        for i in (2..).take_while(|&i| i * i <= r) {
25            if small_p[i] < i {
26                continue;
27            }
28            for j in i..=r0 / i {
29                if small_p[i * j] == i * j {
30                    small_p[i * j] = i;
31                }
32            }
33            for j in (1 + (l - 1) / i..=r / i).filter(|&j| j > 1) {
34                let m = i * j - l;
35                large[m] = false;
36                while large_prod[m] > r0 && large_prod[m] % i == 0 {
37                    large_p[m].push(i);
38                    large_prod[m] /= i;
39                }
40            }
41        }
42        Self { ofs: l, small_p, large_p, large_prod }
43    }
44
45    pub fn factors_dup(&self, n: usize) -> impl Iterator<Item = usize> {
46        let mut res = self.large_p[n - self.ofs].clone();
47        let mut m = self.large_prod[n - self.ofs];
48        if m < self.small_p.len() {
49            while m > 1 {
50                res.push(self.small_p[m]);
51                m /= self.small_p[m];
52            }
53        } else {
54            res.push(m);
55        }
56        res.into_iter()
57    }
58
59    pub fn factors(&self, n: usize) -> impl Iterator<Item = (usize, u32)> {
60        let dup: Vec<_> = self.factors_dup(n).collect();
61        if dup.is_empty() {
62            return vec![].into_iter();
63        }
64        let mut res = vec![(dup[0], 1)];
65        for &p in &dup[1..] {
66            if res.last().unwrap().0 == p {
67                res.last_mut().unwrap().1 += 1;
68            } else {
69                res.push((p, 1));
70            }
71        }
72        res.into_iter()
73    }
74
75    pub fn divisors(&self, n: usize) -> impl Iterator<Item = usize> {
76        let mut res = vec![1];
77        for (p, e) in self.factors(n) {
78            let mut tmp = vec![];
79            let mut pp = 1;
80            for _ in 1..=e {
81                pp *= p;
82                tmp.extend(res.iter().map(|&x| x * pp));
83            }
84            res.extend(tmp);
85        }
86        res.sort_unstable();
87        res.into_iter()
88    }
89}
nekolib/math/segmented_factor_sieve.rs

nekolib/math/
segmented_factor_sieve.rs
