use cargo_snippet::snippet;
#[snippet("sieve")]
pub struct Sieve {
size: usize,
is_prime: Vec<bool>,
primes: Vec<usize>,
}
#[snippet("sieve")]
impl Sieve {
pub fn new(n: usize) -> Self {
let mut spf = vec![None; n + 1];
let mut is_prime = vec![true; n + 1];
let mut primes = Vec::new();
is_prime[0] = false;
is_prime[1] = false;
for i in 2..n + 1 {
if is_prime[i] {
primes.push(i);
spf[i] = Some(i);
}
for prime in &primes {
if i * prime >= n + 1 || prime > &spf[i].unwrap() {
break;
}
is_prime[i * prime] = false;
spf[i * prime] = Some(*prime);
}
}
Self {
size: n,
is_prime,
primes,
}
}
pub fn size(&self) -> usize {
self.size
}
pub fn primes(&self, n: usize) -> Vec<usize> {
assert!(self.size() >= n);
self.primes
.iter()
.take_while(|x| **x <= n)
.cloned()
.collect()
}
pub fn is_prime(&self, n: usize) -> bool {
self.is_prime[n]
}
}
#[snippet("sieve")]
pub fn factorizations_with_sieve(sieve: &Sieve, mut n: usize) -> Vec<(usize, usize)> {
assert!(sieve.size.pow(2) >= n);
let mut res = Vec::new();
let ps = sieve.primes((n as f64).sqrt().ceil() as usize);
for p in ps {
let mut c = 0usize;
while n % p == 0 {
n /= p;
c += 1;
}
if c != 0 {
res.push((p, c));
}
}
if n > 1 {
res.push((n, 1));
}
res
}
#[test]
fn sieve_test() {
let sieve = Sieve::new(10);
assert_eq!(vec![2, 3, 5, 7], sieve.primes(10));
assert_eq!(sieve.is_prime[2], true);
assert_eq!(sieve.is_prime[9], false);
}
#[test]
fn fact_test() {
let sieve = Sieve::new(100000);
let mut f = factorizations_with_sieve(&sieve, 12);
f.sort();
assert_eq!(f, vec![(2, 2), (3, 1)]);
let mut f = factorizations_with_sieve(&sieve, 107);
f.sort();
assert_eq!(f, vec![(107, 1)]);
let mut f = factorizations_with_sieve(&sieve, 1000000007);
f.sort();
assert_eq!(f, vec![(1000000007, 1)]);
}